remove extra stuff

diegoimbert · diegoimbert · commit 3a58cb85983d · 2025-05-09T18:46:34.000+02:00
diff --git a/backend/windmill-common/src/more_serde.rs b/backend/windmill-common/src/more_serde.rs
@@ -8,20 +8,10 @@
 
 //! helpers for serde + serde derive attributes
 
-use crate::{
-    error::{self, to_anyhow},
-    utils::rd_string,
-};
-use bytes::Bytes;
-use futures::TryStreamExt;
-use serde::{de::DeserializeSeed, Deserialize, Deserializer};
-use serde_json::{value::RawValue, Value};
+use crate::utils::rd_string;
+use serde::{Deserialize, Deserializer};
+use serde_json::value::RawValue;
 use std::{fmt::Display, str::FromStr};
-use tokio::{
-    sync::mpsc::Sender,
-    task::{self},
-};
-use tokio_stream::StreamExt;
 
 pub fn default_true() -> bool {
     true
@@ -75,61 +65,3 @@ where
         NumericOrNull::Null => Ok(None),
     }
 }
-
-// Takes a json stream and returns a stream of json values, without loading the
-// entire input into memory.
-pub async fn json_stream_values<
-    'a,
-    D: DeserializeSeed<'a> + 'static + Send,
-    F: FnOnce(Sender<Value>) -> D,
-    E: Display,
->(
-    mut stream: impl TryStreamExt<Item = Result<Bytes, E>> + Send + Unpin + 'static,
-    mpsc_deserializer_factory: F,
-) -> error::Result<impl StreamExt<Item = serde_json::Value>> {
-    const MAX_MPSC_SIZE: usize = 1000;
-
-    use std::path::PathBuf;
-    use tokio::io::AsyncWriteExt;
-
-    let tmp_filename = format!("tmp_json_stream_{}", rd_string(8));
-
-    // Start by writing the async stream (from network) to a file.
-    let mut path = PathBuf::from(std::env::temp_dir());
-    path.push(tmp_filename);
-    let mut file: tokio::fs::File = tokio::fs::File::create(&path).await.map_err(to_anyhow)?;
-    while let Some(chunk) = stream.next().await {
-        let chunk: Bytes = match chunk {
-            Ok(chunk) => chunk,
-            Err(e) => {
-                std::fs::remove_file(&path)?;
-                return Err(error::Error::ExecutionErr(format!(
-                    "Error reading stream: {}",
-                    e
-                )));
-            }
-        };
-        file.write_all(&chunk).await?;
-    }
-
-    let (tx, rx) = tokio::sync::mpsc::channel(MAX_MPSC_SIZE);
-    // Takes ownership of tx
-    let mpsc_deserializer = mpsc_deserializer_factory(tx);
-
-    // We read the file and pipe each element to the channel in a blocking task.
-    let _ = task::spawn_blocking::<_, anyhow::Result<()>>(move || {
-        let sync_file = std::fs::File::open(&path).map_err(to_anyhow)?;
-        let mut buf_reader = std::io::BufReader::new(sync_file);
-
-        let mut deserializer = serde_json::Deserializer::from_reader(&mut buf_reader);
-        // This deserializer will read the file and send each row to the channel
-        let _ = mpsc_deserializer.deserialize(&mut deserializer)?;
-
-        std::fs::remove_file(&path)?;
-        Ok(())
-        // tx drops with mpsc_deserializer so the stream ends
-    });
-
-    let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
-    Ok(stream)
-}
diff --git a/backend/windmill-worker/src/snowflake_executor.rs b/backend/windmill-worker/src/snowflake_executor.rs
@@ -2,16 +2,13 @@ use base64::{engine, Engine as _};
 use chrono::Datelike;
 use core::fmt::Write;
 use futures::future::BoxFuture;
-use futures::{FutureExt, StreamExt, TryFutureExt};
+use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt};
 use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
 use reqwest::{Client, Response};
 use serde_json::{json, value::RawValue, Value};
 use sha2::{Digest, Sha256};
-use std::collections::{BTreeMap, HashMap};
-use std::convert::Infallible;
+use std::collections::HashMap;
 use windmill_common::error::to_anyhow;
-use windmill_common::more_serde::json_stream_values;
-use windmill_common::s3_helpers::convert_json_line_stream;
 use windmill_common::worker::Connection;
 
 use windmill_common::{error::Error, worker::to_raw_value};
@@ -179,30 +176,12 @@ fn do_snowflake_inner<'a>(
         if skip_collect {
             handle_snowflake_result(result).await?;
             Ok(to_raw_value(&Value::Array(vec![])))
-        } else if let Some(ref s3) = s3 {
-            // do not do parse_snowflake_response as it will call .json() and
-            // load the entire response into memory
-            let result = result.map_err(|e| {
-                Error::ExecutionErr(format!("Could not send request to Snowflake: {:?}", e))
-            })?;
-
-            let rows_stream = json_stream_values(result.bytes_stream(), |sender| {
-                RootMpscDeserializer { sender }
-            })
-            .await?
-            .boxed()
-            .map(|chunk| Ok::<_, Infallible>(chunk));
-
-            let stream = convert_json_line_stream(rows_stream, s3.format).await?;
-            s3.upload(stream.boxed()).await?;
-
-            Ok(serde_json::value::to_raw_value(&s3.object_key)?)
         } else {
             let response = result
                 .parse_snowflake_response::<SnowflakeResponse>()
                 .await?;
 
-            if response.resultSetMetaData.numRows > 10000 {
+            if s3.is_none() && response.resultSetMetaData.numRows > 10000 {
                 return Err(Error::ExecutionErr(
                     "More than 10000 rows were requested, use LIMIT 10000 to limit the number of rows"
                         .to_string(),
@@ -219,54 +198,66 @@ fn do_snowflake_inner<'a>(
                 );
             }
 
-            let mut rows = response.data;
-
-            if response.resultSetMetaData.partitionInfo.len() > 1 {
-                for idx in 1..response.resultSetMetaData.partitionInfo.len() {
-                    let url = format!(
-                        "https://{}.snowflakecomputing.com/api/v2/statements/{}",
-                        account_identifier.to_uppercase(),
-                        response.statementHandle
-                    );
-                    let mut request = HTTP_CLIENT
-                        .get(url)
-                        .bearer_auth(token)
-                        .query(&[("partition", idx.to_string())]);
-
-                    if token_is_keypair {
-                        request =
-                            request.header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT");
-                    }
-
-                    let response = request
-                        .send()
-                        .await
-                        .parse_snowflake_response::<SnowflakeDataOnlyResponse>()
-                        .await?;
+            let rows_stream = async_stream::stream! {
+                for row in response.data {
+                    yield Ok::<Vec<Value>, windmill_common::error::Error>(row);
+                }
 
-                    rows.extend(response.data);
+                if response.resultSetMetaData.partitionInfo.len() > 1 {
+                    for idx in 1..response.resultSetMetaData.partitionInfo.len() {
+                        let url = format!(
+                            "https://{}.snowflakecomputing.com/api/v2/statements/{}",
+                            account_identifier.to_uppercase(),
+                            response.statementHandle
+                        );
+                        let mut request = HTTP_CLIENT
+                            .get(url)
+                            .bearer_auth(token)
+                            .query(&[("partition", idx.to_string())]);
+
+                        if token_is_keypair {
+                            request =
+                                request.header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT");
+                        }
+
+                        let response = request
+                            .send()
+                            .await
+                            .parse_snowflake_response::<SnowflakeDataOnlyResponse>()
+                            .await?;
+
+                        for row in response.data {
+                            yield Ok(row);
+                        }
+                    }
                 }
+            };
+
+            let rows_stream = rows_stream.map_ok(|row| {
+                let mut row_map = serde_json::Map::new();
+                row.iter()
+                    .zip(response.resultSetMetaData.rowType.iter())
+                    .for_each(|(val, row_type)| {
+                        row_map.insert(row_type.name.clone(), parse_val(&val, &row_type.r#type));
+                    });
+                row_map
+            });
+
+            if let Some(s3) = s3 {
+                // let rows_stream =
+                //     rows_stream.map(|r| serde_json::value::to_value(&r?).map_err(to_anyhow));
+                // let stream = convert_json_line_stream(rows_stream.boxed(), s3.format).await?;
+                // TODO fix this
+                // s3.upload(stream.boxed()).await?;
+                Ok(to_raw_value(&s3.object_key))
+            } else {
+                let rows = rows_stream
+                    .collect::<Vec<_>>()
+                    .await
+                    .into_iter()
+                    .collect::<Result<Vec<_>, _>>()?;
+                Ok(to_raw_value(&rows))
             }
-
-            let rows = to_raw_value(
-                &rows
-                    .iter()
-                    .map(|row| {
-                        let mut row_map = serde_json::Map::new();
-                        row.iter()
-                            .zip(response.resultSetMetaData.rowType.iter())
-                            .for_each(|(val, row_type)| {
-                                row_map.insert(
-                                    row_type.name.clone(),
-                                    parse_val(&val, &row_type.r#type),
-                                );
-                            });
-                        row_map
-                    })
-                    .collect::<Vec<_>>(),
-            );
-
-            Ok(rows)
         }
     };
 
@@ -587,103 +578,3 @@ fn parse_val(value: &Value, typ: &str) -> Value {
         ))
     }
 }
-
-// This deserializer takes a snowflake response as a stream and sends each row to an mpsc
-// channel as a json record without storing the full input json in memory.
-struct RootMpscDeserializer {
-    sender: tokio::sync::mpsc::Sender<serde_json::Value>,
-}
-
-impl<'de> serde::de::DeserializeSeed<'de> for RootMpscDeserializer {
-    type Value = ();
-    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        struct RootVisitor<'a> {
-            sender: &'a tokio::sync::mpsc::Sender<serde_json::Value>,
-            col_defs: Vec<String>,
-        }
-
-        impl<'de, 'a> serde::de::Visitor<'de> for RootVisitor<'a> {
-            type Value = ();
-            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
-                formatter.write_str("data field from snowflake response")
-            }
-            fn visit_map<A>(mut self, mut map: A) -> Result<(), A::Error>
-            where
-                A: serde::de::MapAccess<'de>,
-            {
-                while let Some(key) = map.next_key::<String>()? {
-                    if key == "resultSetMetaData" {
-                        let result_set_metadata: SnowflakeResultSetMetaData = map.next_value()?;
-                        self.col_defs = result_set_metadata
-                            .rowType
-                            .iter()
-                            .map(|x| x.name.clone())
-                            .collect::<Vec<String>>();
-                    } else if key == "data" {
-                        let () = map.next_value_seed(RowsMpscDeserializer {
-                            sender: self.sender,
-                            col_defs: &self.col_defs,
-                        })?;
-                    } else {
-                        let _: serde::de::IgnoredAny = map.next_value()?;
-                    }
-                }
-                Ok(())
-            }
-        }
-
-        deserializer.deserialize_map(RootVisitor { sender: &self.sender, col_defs: vec![] })
-    }
-}
-
-struct RowsMpscDeserializer<'a> {
-    sender: &'a tokio::sync::mpsc::Sender<serde_json::Value>,
-    col_defs: &'a Vec<String>,
-}
-
-impl<'de, 'a> serde::de::DeserializeSeed<'de> for RowsMpscDeserializer<'a> {
-    type Value = ();
-    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        struct RowsVisitor<'a> {
-            sender: &'a tokio::sync::mpsc::Sender<serde_json::Value>,
-            col_defs: &'a Vec<String>,
-        }
-
-        impl<'de, 'a> serde::de::Visitor<'de> for RowsVisitor<'a> {
-            type Value = ();
-
-            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
-                formatter.write_str("a sequence of rows")
-            }
-
-            fn visit_seq<A>(self, mut seq: A) -> Result<(), A::Error>
-            where
-                A: serde::de::SeqAccess<'de>,
-            {
-                while let Some(elem) = seq.next_element::<Vec<Value>>()? {
-                    let mut row = BTreeMap::<&str, Value>::new();
-                    for (i, field) in elem.iter().enumerate() {
-                        let col_name = self.col_defs.get(i).map(|s| s.as_str()).unwrap_or("");
-                        row.insert(col_name, field.clone());
-                    }
-                    let row = serde_json::to_value(row).map_err(|err| {
-                        serde::de::Error::custom(format!("Map parse failed: {err}"))
-                    })?;
-                    self.sender.blocking_send(row).map_err(|err| {
-                        serde::de::Error::custom(format!("Queue send failed: {err}"))
-                    })?;
-                }
-
-                Ok(())
-            }
-        }
-
-        deserializer.deserialize_seq(RowsVisitor { sender: &self.sender, col_defs: &self.col_defs })
-    }
-}
