goat_cli/count/
mod.rs

1//!
2//! Invoked by calling:
3//! `goat-cli count <args>`
4
5use anyhow::{bail, Result};
6use futures::StreamExt;
7use reqwest;
8use reqwest::header::ACCEPT;
9use serde_json::Value;
10
11use crate::utils::cli_matches;
12use crate::IndexType;
13
14/// `goat-cli count` presents an identical CLI to `goat-cli search` but prints
15/// to the console in the CLI call here, and to the stderr in the `goat-cli search` call.
16pub async fn count(
17    matches: &clap::ArgMatches,
18    cli: bool,
19    print_warning: bool,
20    unique_ids: Vec<String>,
21    index_type: IndexType,
22) -> Result<Option<u64>> {
23    let (size_int, url_vector, url_vector_api) =
24        cli_matches::process_cli_args(matches, "count", unique_ids, index_type)?;
25    let concurrent_requests = url_vector_api.len();
26
27    let fetches = futures::stream::iter(url_vector_api.iter().map(|path| async move {
28        // possibly make a again::RetryPolicy
29        // to catch all the values in a *very* large request.
30        let client = reqwest::Client::new();
31
32        match again::retry(|| client.get(path).header(ACCEPT, "application/json").send()).await {
33            Ok(resp) => match resp.text().await {
34                Ok(body) => {
35                    let v: Value = serde_json::from_str(&body)?;
36                    let count = &v["count"].as_u64();
37                    match count {
38                        Some(c) => Ok(*c),
39                        None => Ok(0), // bail!("Bad response."),
40                    }
41                }
42                Err(_) => bail!("ERROR reading {}", path),
43            },
44            Err(_) => bail!("ERROR downloading {}", path),
45        }
46    }))
47    .buffered(concurrent_requests)
48    .collect::<Vec<_>>();
49
50    let awaited_fetches = fetches.await;
51
52    match cli {
53        true => {
54            // print to console
55            let mut outer_count = 0;
56            println!("search_query\tcount");
57            for (el1, el2) in awaited_fetches.iter().zip(url_vector.iter()) {
58                let count = match el1 {
59                    Ok(e) => e,
60                    Err(e) => bail!("{}", e),
61                };
62                println!("{}\t{}", el2, count);
63                outer_count += *count;
64            }
65            Ok(Some(outer_count))
66        }
67        false => {
68            // need
69            let mut outer_count = 0;
70            // the zip does not correspond to the awaited fetches...
71            // need to match them
72            for (el1, el2) in awaited_fetches.iter().zip(url_vector.iter()) {
73                let count = match el1 {
74                    Ok(e) => e,
75                    Err(e) => bail!("{}", e),
76                };
77                if print_warning && size_int < *count {
78                    eprintln!("For search query {}, size specified ({}) was less than the number of results returned, ({}).", el2, size_int, count)
79                }
80                outer_count += *count;
81            }
82
83            Ok(Some(outer_count))
84        }
85    }
86}
goat_cli/count/mod.rs

goat_cli/count/
mod.rs
