extern crate firn;
use anyhow::Result;
use arrow::util::pretty::pretty_format_batches;
use firn::{QueryData, SnowflakeApi};
const RECURSIVE_CTE: &str = "\
WITH RECURSIVE Compute_CTE AS (
SELECT 1 AS n UNION ALL
SELECT n + 1 FROM Compute_CTE WHERE n < 10000000
)
SELECT SUM(n) FROM Compute_CTE";
#[tokio::main]
async fn main() -> Result<()> {
pretty_env_logger::init();
let api = SnowflakeApi::from_env()?;
let sql = std::env::var("SQL").ok();
let sql = sql.as_deref().unwrap_or(RECURSIVE_CTE);
log::info!("Running long query (this should hit the async path)");
let started = std::time::Instant::now();
let res = api.exec(sql).await?;
log::info!("Query finished in {:?}", started.elapsed());
log::info!("query_id: {}", res.metadata.query_id);
match res.data {
QueryData::Arrow(a) => println!("{}", pretty_format_batches(&a).unwrap()),
QueryData::Json(j) => println!("{j}"),
QueryData::Empty => println!("Query finished successfully"),
}
Ok(())
}