use serde::Deserialize;
use crate::clouds::gcloud::bigquery;
use crate::common::*;
use crate::drivers::{
bigquery::BigQueryLocator,
bigquery_shared::{BqTable, GCloudDriverArguments, Usage},
};
#[instrument(
level = "trace",
name = "bigquery::count",
skip(shared_args, source_args)
)]
pub(crate) async fn count_helper(
locator: BigQueryLocator,
shared_args: SharedArguments<Unverified>,
source_args: SourceArguments<Unverified>,
) -> Result<usize> {
let shared_args = shared_args.verify(BigQueryLocator::features())?;
let source_args = source_args.verify(BigQueryLocator::features())?;
let job_labels = source_args
.driver_args()
.deserialize::<GCloudDriverArguments>()
.context("error parsing --from-args")?
.job_labels
.to_owned();
let schema = shared_args.schema();
let table_name = locator.as_table_name().to_owned();
let table = BqTable::for_table_name_and_columns(
schema,
table_name,
&schema.table.columns,
Usage::FinalTable,
)?;
let mut count_sql_data = vec![];
table.write_count_sql(&source_args, &mut count_sql_data)?;
let count_sql = String::from_utf8(count_sql_data).expect("should always be UTF-8");
debug!("count SQL: {}", count_sql);
#[derive(Deserialize)]
struct CountRow {
count: String,
}
let count_str =
bigquery::query_one::<CountRow>(locator.project(), &count_sql, &job_labels)
.await?
.count;
count_str
.parse::<usize>()
.context("could not parse count output")
}