S3 works with new syntax

Author: diegoimbert

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

@@ -121,14 +121,14 @@ pub fn parse_db_resource(code: &str) -> Option<String> {
 }
 
 pub struct S3ModeArgs {
-    pub object_key: String,
+    pub folder_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 })
+    Some(S3ModeArgs { folder_key: arg1, storage: arg2 })
 }
 
 pub fn parse_sql_blocks(code: &str) -> Vec<&str> {

backend/windmill-worker/src/bigquery_executor.rs

@@ -66,10 +66,12 @@ struct BigqueryError {
     message: String,
 }
 
-struct S3Mode<'a> {
-    client: &'a AuthedClient,
+#[derive(Clone)]
+struct S3Mode {
+    client: AuthedClient,
     object_key: String,
     storage: Option<String>,
+    workspace_id: String,
 }
 
 fn do_bigquery_inner<'a>(
@@ -81,7 +83,7 @@ fn do_bigquery_inner<'a>(
     column_order: Option<&'a mut Option<Vec<String>>>,
     skip_collect: bool,
     http_client: &'a Client,
-    s3: &Option<S3Mode>,
+    s3: Option<S3Mode>,
 ) -> windmill_common::error::Result<BoxFuture<'a, windmill_common::error::Result<Box<RawValue>>>> {
     let param_names = parse_sql_statement_named_params(query, '@');
 
@@ -122,7 +124,16 @@ fn do_bigquery_inner<'a>(
                 if skip_collect {
                     return Ok(to_raw_value(&Value::Array(vec![])));
                 } else if let Some(ref s3) = s3 {
-                    // s3.client.get("", query)
+                    s3.client
+                        .upload_s3_file(
+                            s3.workspace_id.as_str(),
+                            s3.object_key.clone(),
+                            s3.storage.clone(),
+                            response.bytes_stream(),
+                        )
+                        .await?;
+
+                    Ok(serde_json::value::to_raw_value(&s3.object_key)?)
                 } else {
                     let result = response.json::<BigqueryResponse>().await.map_err(|e| {
                         Error::ExecutionErr(format!(
@@ -231,9 +242,10 @@ pub async fn do_bigquery(
 
     let inline_db_res_path = parse_db_resource(&query);
     let s3 = parse_s3_mode(&query).map(|s3_mode| S3Mode {
-        client: &client,
+        client: client.clone(),
         storage: s3_mode.storage,
-        object_key: s3_mode.object_key,
+        object_key: format!("{}/{}.txt", s3_mode.folder_key, job.id),
+        workspace_id: job.workspace_id.clone(),
     });
 
     let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
@@ -347,7 +359,7 @@ pub async fn do_bigquery(
                     None,
                     annotations.return_last_result && i < queries.len() - 1,
                     &http_client,
-                    &s3,
+                    s3.clone(),
                 )
             })
             .collect::<windmill_common::error::Result<Vec<_>>>()?;
@@ -377,7 +389,7 @@ pub async fn do_bigquery(
             Some(column_order),
             false,
             &http_client,
-            &s3,
+            s3,
         )?
     };
 

backend/windmill-worker/src/worker.rs

@@ -39,7 +39,7 @@ use windmill_common::METRICS_DEBUG_ENABLED;
 #[cfg(feature = "prometheus")]
 use windmill_common::METRICS_ENABLED;
 
-use reqwest::Response;
+use reqwest::{Body, Response};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use sqlx::types::Json;
 use std::{
@@ -520,6 +520,44 @@ impl AuthedClient {
             _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default())),
         }
     }
+
+    pub async fn upload_s3_file<S>(
+        &self,
+        workspace_id: &str,
+        object_key: String,
+        storage: Option<String>,
+        body: S,
+    ) -> anyhow::Result<Response>
+    where
+        S: futures::stream::TryStream + Send + 'static,
+        S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
+        bytes::Bytes: From<S::Ok>,
+    {
+        let mut query = vec![("file_key", object_key)];
+        if let Some(storage) = storage {
+            query.push(("storage", storage.clone()));
+        }
+        self.force_client
+            .as_ref()
+            .unwrap_or(&HTTP_CLIENT)
+            .post(format!(
+                "{}/api/w/{}/job_helpers/upload_s3_file",
+                self.base_internal_url, workspace_id
+            ))
+            .query(&query)
+            .header(
+                reqwest::header::ACCEPT,
+                reqwest::header::HeaderValue::from_static("application/json"),
+            )
+            .header(
+                reqwest::header::AUTHORIZATION,
+                reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))?,
+            )
+            .body(Body::wrap_stream(body))
+            .send()
+            .await
+            .context(format!("Sent upload_s3_file request",))
+    }
 }
 
 #[derive(Clone)]