reinfer-cli 0.38.5

Command line interface for Re:infer, the conversational data intelligence platform
use anyhow::{Context, Result};

use colored::Colorize;
use reinfer_client::{resources::bucket_statistics::Count, BucketIdentifier, Client, EmailId};
use std::{
    fs::File,
    io::{self, BufWriter, Write},
    path::PathBuf,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
};
use structopt::StructOpt;

use crate::{
    printer::print_resources_as_json,
    progress::{Options as ProgressOptions, Progress},
};

#[derive(Debug, StructOpt)]
pub struct GetManyEmailsArgs {
    #[structopt(name = "bucket")]
    /// Bucket name or id
    bucket: BucketIdentifier,

    #[structopt(short = "f", long = "file", parse(from_os_str))]
    /// Path where to write comments as JSON. If not specified, stdout will be used.
    path: Option<PathBuf>,

    #[structopt(name = "id")]
    /// Id of specific email to return
    id: Option<EmailId>,
}

pub fn get_many(client: &Client, args: &GetManyEmailsArgs) -> Result<()> {
    let GetManyEmailsArgs { bucket, path, id } = args;

    let file = match path {
        Some(path) => Some(
            File::create(path)
                .with_context(|| format!("Could not open file for writing `{}`", path.display()))
                .map(BufWriter::new)?,
        ),
        None => None,
    };

    if let Some(id) = id {
        if let Some(file) = file {
            return download_email(client, bucket.clone(), id.clone(), file);
        } else {
            return download_email(client, bucket.clone(), id.clone(), io::stdout().lock());
        }
    }

    if let Some(file) = file {
        download_emails(client, bucket.clone(), file)
    } else {
        download_emails(client, bucket.clone(), io::stdout().lock())
    }
}

fn download_email(
    client: &Client,
    bucket_identifier: BucketIdentifier,
    id: EmailId,
    mut writer: impl Write,
) -> Result<()> {
    let bucket = client
        .get_bucket(bucket_identifier)
        .context("Operation to get bucket has failed.")?;

    let response = client.get_email(&bucket.full_name(), id)?;

    print_resources_as_json(response, &mut writer)
}

fn download_emails(
    client: &Client,
    bucket_identifier: BucketIdentifier,
    mut writer: impl Write,
) -> Result<()> {
    let bucket = client
        .get_bucket(bucket_identifier)
        .context("Operation to get bucket has failed.")?;

    let bucket_statistics = client
        .get_bucket_statistics(&bucket.full_name())
        .context("Could not get bucket statistics")?;

    let statistics = Arc::new(Statistics::new());

    let progress_bytes = match bucket_statistics.count {
        Count::LowerBoundBucketCount { value } => value,
        Count::ExactBucketCount { value } => value,
    } as u64;

    let _progress = get_emails_progress_bar(progress_bytes, &statistics);

    client
        .get_emails_iter(&bucket.full_name(), None)
        .try_for_each(|page| {
            let page = page.context("Operation to get emails has failed.")?;
            statistics.add_emails(page.len());
            print_resources_as_json(page.into_iter(), &mut writer)
        })?;
    log::info!(
        "Successfully downloaded {} emails.",
        statistics.num_downloaded(),
    );
    Ok(())
}

#[derive(Debug)]
pub struct Statistics {
    downloaded: AtomicUsize,
}

impl Statistics {
    fn new() -> Self {
        Self {
            downloaded: AtomicUsize::new(0),
        }
    }

    #[inline]
    fn add_emails(&self, num_downloaded: usize) {
        self.downloaded.fetch_add(num_downloaded, Ordering::SeqCst);
    }

    #[inline]
    fn num_downloaded(&self) -> usize {
        self.downloaded.load(Ordering::SeqCst)
    }
}

fn get_emails_progress_bar(total_bytes: u64, statistics: &Arc<Statistics>) -> Progress {
    Progress::new(
        move |statistics| {
            let num_downloaded = statistics.num_downloaded();
            (
                num_downloaded as u64,
                format!(
                    "{} {}",
                    num_downloaded.to_string().bold(),
                    "emails".dimmed(),
                ),
            )
        },
        statistics,
        Some(total_bytes),
        ProgressOptions { bytes_units: false },
    )
}