stream to s3 boilerplate

Author: diegoimbert

backend/Cargo.lock

@@ -120,6 +120,17 @@ pub fn parse_db_resource(code: &str) -> Option<String> {
     cap.map(|x| x.get(1).map(|x| x.as_str().to_string()).unwrap())
 }
 
+pub struct S3ModeArgs {
+    pub object_key: String,
+    pub storage: Option<String>,
+}
+pub fn parse_s3_mode(code: &str) -> Option<S3ModeArgs> {
+    let cap = RE_S3_MODE.captures(code)?;
+    let arg1 = cap.get(1).map(|x| x.as_str().to_string())?;
+    let arg2 = cap.get(2).map(|x| x.as_str().to_string());
+    Some(S3ModeArgs { object_key: arg1, storage: arg2 })
+}
+
 pub fn parse_sql_blocks(code: &str) -> Vec<&str> {
     let mut blocks = vec![];
     let mut last_idx = 0;
@@ -147,6 +158,7 @@ lazy_static::lazy_static! {
     static ref RE_NONEMPTY_SQL_BLOCK: Regex = Regex::new(r#"(?m)^\s*[^\s](?:[^-]|$)"#).unwrap();
 
     static ref RE_DB: Regex = Regex::new(r#"(?m)^-- database (\S+) *(?:\r|\n|$)"#).unwrap();
+    static ref RE_S3_MODE: Regex = Regex::new(r#"(?m)^-- s3 (\S+)( +(\S+))? *(?:\r|\n|$)"#).unwrap();
 
     // -- $1 name (type) = default
     static ref RE_ARG_MYSQL: Regex = Regex::new(r#"(?m)^-- \? (\w+) \((\w+)\)(?: ?\= ?(.+))? *(?:\r|\n|$)"#).unwrap();

backend/parsers/windmill-parser-sql/src/lib.rs

@@ -8,7 +8,8 @@ use windmill_common::error::to_anyhow;
 use windmill_common::worker::Connection;
 use windmill_common::{error::Error, worker::to_raw_value};
 use windmill_parser_sql::{
-    parse_bigquery_sig, parse_db_resource, parse_sql_blocks, parse_sql_statement_named_params,
+    parse_bigquery_sig, parse_db_resource, parse_s3_mode, parse_sql_blocks,
+    parse_sql_statement_named_params,
 };
 use windmill_queue::CanceledBy;
 
@@ -65,6 +66,12 @@ struct BigqueryError {
     message: String,
 }
 
+struct S3Mode<'a> {
+    client: &'a AuthedClient,
+    object_key: String,
+    storage: Option<String>,
+}
+
 fn do_bigquery_inner<'a>(
     query: &'a str,
     all_statement_values: &'a HashMap<String, Value>,
@@ -74,6 +81,7 @@ fn do_bigquery_inner<'a>(
     column_order: Option<&'a mut Option<Vec<String>>>,
     skip_collect: bool,
     http_client: &'a Client,
+    s3: &Option<S3Mode>,
 ) -> windmill_common::error::Result<BoxFuture<'a, windmill_common::error::Result<Box<RawValue>>>> {
     let param_names = parse_sql_statement_named_params(query, '@');
 
@@ -113,6 +121,8 @@ fn do_bigquery_inner<'a>(
             Ok(_) => {
                 if skip_collect {
                     return Ok(to_raw_value(&Value::Array(vec![])));
+                } else if let Some(ref s3) = s3 {
+                    // s3.client.get("", query)
                 } else {
                     let result = response.json::<BigqueryResponse>().await.map_err(|e| {
                         Error::ExecutionErr(format!(
@@ -220,6 +230,11 @@ pub async fn do_bigquery(
     let bigquery_args = build_args_values(job, client, conn).await?;
 
     let inline_db_res_path = parse_db_resource(&query);
+    let s3 = parse_s3_mode(&query).map(|s3_mode| S3Mode {
+        client: &client,
+        storage: s3_mode.storage,
+        object_key: s3_mode.object_key,
+    });
 
     let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
         Some(
@@ -332,6 +347,7 @@ pub async fn do_bigquery(
                     None,
                     annotations.return_last_result && i < queries.len() - 1,
                     &http_client,
+                    &s3,
                 )
             })
             .collect::<windmill_common::error::Result<Vec<_>>>()?;
@@ -361,6 +377,7 @@ pub async fn do_bigquery(
             Some(column_order),
             false,
             &http_client,
+            &s3,
         )?
     };