fix: relabel materialization timeout and document per-phase budget

backend/windmill-worker/src/bigquery_executor.rs

@@ -398,32 +398,49 @@ pub async fn do_bigquery(
     // Materialize any `(s3object)` args into JSON text and rewrite the arg type to
     // STRING. The user wraps the parameter with `JSON_EXTRACT_ARRAY(@p)` (or similar)
     // in their SQL.
-    for arg in sig.iter_mut() {
-        if arg.otyp.as_deref() != Some("s3object") {
-            continue;
-        }
-        let raw = bigquery_args.remove(&arg.name).unwrap_or(Value::Null);
-        if matches!(raw, Value::Null) {
-            return Err(Error::BadRequest(format!(
-                "Missing S3Object value for arg `{}`",
-                arg.name
-            )));
-        }
-        let s3_obj: windmill_types::s3::S3Object = serde_json::from_value(raw).map_err(|e| {
-            Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
-        })?;
-        let json_text =
-            crate::sql_s3_input::fetch_s3object_as_json_text(client, &job.workspace_id, &s3_obj)
-                .await
-                .map_err(|e| {
-                    Error::ExecutionErr(format!(
-                        "Failed to fetch S3 object for arg `{}`: {e}",
-                        arg.name
-                    ))
+    let materialize_f = async {
+        for arg in sig.iter_mut() {
+            if arg.otyp.as_deref() != Some("s3object") {
+                continue;
+            }
+            let raw = bigquery_args.remove(&arg.name).unwrap_or(Value::Null);
+            if matches!(raw, Value::Null) {
+                return Err(Error::BadRequest(format!(
+                    "Missing S3Object value for arg `{}`",
+                    arg.name
+                )));
+            }
+            let s3_obj: windmill_types::s3::S3Object =
+                serde_json::from_value(raw).map_err(|e| {
+                    Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
                 })?;
-        bigquery_args.insert(arg.name.clone(), Value::String(json_text));
-        arg.otyp = Some("string".to_string());
-    }
+            let json_text = crate::sql_s3_input::fetch_s3object_as_json_text(
+                client,
+                &job.workspace_id,
+                &s3_obj,
+            )
+            .await
+            .map_err(|e| {
+                Error::ExecutionErr(format!(
+                    "Failed to fetch S3 object for arg `{}`: {e}",
+                    arg.name
+                ))
+            })?;
+            bigquery_args.insert(arg.name.clone(), Value::String(json_text));
+            arg.otyp = Some("string".to_string());
+        }
+        Ok(())
+    };
+    crate::sql_s3_input::materialize_under_job_poller(
+        job,
+        conn,
+        mem_peak,
+        canceled_by,
+        worker_name,
+        occupancy_metrics,
+        materialize_f,
+    )
+    .await?;
 
     let reserved_variables =
         get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;

backend/windmill-worker/src/mssql_executor.rs

@@ -223,30 +223,43 @@ pub async fn do_mssql(
 
     // Materialize any `(s3object)` args into JSON text. tiberius binds the resulting
     // String as nvarchar(max), which is exactly the input type for `OPENJSON(@P)`.
-    for arg in sig.iter() {
-        if arg.otyp.as_deref() != Some("s3object") {
-            continue;
-        }
-        let raw = mssql_args.remove(&arg.name).unwrap_or(Value::Null);
-        if matches!(raw, Value::Null) {
-            return Err(Error::BadRequest(format!(
-                "Missing S3Object value for arg `{}`",
-                arg.name
-            )));
-        }
-        let s3_obj: S3Object = serde_json::from_value(raw).map_err(|e| {
-            Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
-        })?;
-        let json_text = fetch_s3object_as_json_text(authed_client, &job.workspace_id, &s3_obj)
-            .await
-            .map_err(|e| {
-                Error::ExecutionErr(format!(
-                    "Failed to fetch S3 object for arg `{}`: {e}",
+    let materialize_f = async {
+        for arg in sig.iter() {
+            if arg.otyp.as_deref() != Some("s3object") {
+                continue;
+            }
+            let raw = mssql_args.remove(&arg.name).unwrap_or(Value::Null);
+            if matches!(raw, Value::Null) {
+                return Err(Error::BadRequest(format!(
+                    "Missing S3Object value for arg `{}`",
                     arg.name
-                ))
+                )));
+            }
+            let s3_obj: S3Object = serde_json::from_value(raw).map_err(|e| {
+                Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
             })?;
-        mssql_args.insert(arg.name.clone(), Value::String(json_text));
-    }
+            let json_text = fetch_s3object_as_json_text(authed_client, &job.workspace_id, &s3_obj)
+                .await
+                .map_err(|e| {
+                    Error::ExecutionErr(format!(
+                        "Failed to fetch S3 object for arg `{}`: {e}",
+                        arg.name
+                    ))
+                })?;
+            mssql_args.insert(arg.name.clone(), Value::String(json_text));
+        }
+        Ok(())
+    };
+    crate::sql_s3_input::materialize_under_job_poller(
+        job,
+        conn,
+        mem_peak,
+        canceled_by,
+        worker_name,
+        occupancy_metrics,
+        materialize_f,
+    )
+    .await?;
 
     let reserved_variables =
         get_reserved_variables(job, &authed_client.token, conn, parent_runnable_path).await?;

backend/windmill-worker/src/mysql_executor.rs

@@ -228,31 +228,48 @@ pub async fn do_mysql(
 
     // Materialize any `(s3object)` args into JSON text. mysql_async binds strings as
     // VARBINARY/TEXT, which MySQL's `JSON_TABLE`/`JSON_EXTRACT` accept directly.
-    for arg in sig.iter() {
-        if arg.otyp.as_deref() != Some("s3object") {
-            continue;
-        }
-        let raw = job_args.remove(&arg.name).unwrap_or(Value::Null);
-        if matches!(raw, Value::Null) {
-            return Err(Error::BadRequest(format!(
-                "Missing S3Object value for arg `{}`",
-                arg.name
-            )));
-        }
-        let s3_obj: windmill_types::s3::S3Object = serde_json::from_value(raw).map_err(|e| {
-            Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
-        })?;
-        let json_text =
-            crate::sql_s3_input::fetch_s3object_as_json_text(client, &job.workspace_id, &s3_obj)
-                .await
-                .map_err(|e| {
-                    Error::ExecutionErr(format!(
-                        "Failed to fetch S3 object for arg `{}`: {e}",
-                        arg.name
-                    ))
+    let materialize_f = async {
+        for arg in sig.iter() {
+            if arg.otyp.as_deref() != Some("s3object") {
+                continue;
+            }
+            let raw = job_args.remove(&arg.name).unwrap_or(Value::Null);
+            if matches!(raw, Value::Null) {
+                return Err(Error::BadRequest(format!(
+                    "Missing S3Object value for arg `{}`",
+                    arg.name
+                )));
+            }
+            let s3_obj: windmill_types::s3::S3Object =
+                serde_json::from_value(raw).map_err(|e| {
+                    Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
                 })?;
-        job_args.insert(arg.name.clone(), Value::String(json_text));
-    }
+            let json_text = crate::sql_s3_input::fetch_s3object_as_json_text(
+                client,
+                &job.workspace_id,
+                &s3_obj,
+            )
+            .await
+            .map_err(|e| {
+                Error::ExecutionErr(format!(
+                    "Failed to fetch S3 object for arg `{}`: {e}",
+                    arg.name
+                ))
+            })?;
+            job_args.insert(arg.name.clone(), Value::String(json_text));
+        }
+        Ok(())
+    };
+    crate::sql_s3_input::materialize_under_job_poller(
+        job,
+        conn,
+        mem_peak,
+        canceled_by,
+        worker_name,
+        occupancy_metrics,
+        materialize_f,
+    )
+    .await?;
 
     let reserved_variables =
         get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;

backend/windmill-worker/src/pg_executor.rs

@@ -768,7 +768,22 @@ pub async fn do_postgresql(
     // Materialize any `(s3object)` args into JSON text and rebind them as `jsonb` so
     // `otyp_to_pg_type` picks the right binding. Must run before the param map is
     // built below.
-    materialize_s3object_args(&mut sig.args, &mut pg_args, client, &job.workspace_id).await?;
+    let materialize_f =
+        materialize_s3object_args(&mut sig.args, &mut pg_args, client, &job.workspace_id);
+    if run_inline {
+        materialize_f.await?;
+    } else {
+        crate::sql_s3_input::materialize_under_job_poller(
+            job,
+            conn,
+            mem_peak,
+            canceled_by,
+            worker_name,
+            occupancy_metrics,
+            materialize_f,
+        )
+        .await?;
+    }
 
     let reserved_variables =
         get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;

backend/windmill-worker/src/snowflake_executor.rs

@@ -550,35 +550,48 @@ pub async fn do_snowflake(
         let sig = parse_snowflake_sig(query)
             .map_err(|x| Error::ExecutionErr(x.to_string()))?
             .args;
-        for arg in sig.iter() {
-            if arg.otyp.as_deref() != Some("s3object") {
-                continue;
-            }
-            let raw = snowflake_args.remove(&arg.name).unwrap_or(Value::Null);
-            if matches!(raw, Value::Null) {
-                return Err(Error::BadRequest(format!(
-                    "Missing S3Object value for arg `{}`",
-                    arg.name
-                )));
-            }
-            let s3_obj: windmill_types::s3::S3Object =
-                serde_json::from_value(raw).map_err(|e| {
-                    Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
+        let materialize_f = async {
+            for arg in sig.iter() {
+                if arg.otyp.as_deref() != Some("s3object") {
+                    continue;
+                }
+                let raw = snowflake_args.remove(&arg.name).unwrap_or(Value::Null);
+                if matches!(raw, Value::Null) {
+                    return Err(Error::BadRequest(format!(
+                        "Missing S3Object value for arg `{}`",
+                        arg.name
+                    )));
+                }
+                let s3_obj: windmill_types::s3::S3Object =
+                    serde_json::from_value(raw).map_err(|e| {
+                        Error::ExecutionErr(format!("Invalid S3Object for arg `{}`: {e}", arg.name))
+                    })?;
+                let json_text = crate::sql_s3_input::fetch_s3object_as_json_text(
+                    client,
+                    &job.workspace_id,
+                    &s3_obj,
+                )
+                .await
+                .map_err(|e| {
+                    Error::ExecutionErr(format!(
+                        "Failed to fetch S3 object for arg `{}`: {e}",
+                        arg.name
+                    ))
                 })?;
-            let json_text = crate::sql_s3_input::fetch_s3object_as_json_text(
-                client,
-                &job.workspace_id,
-                &s3_obj,
-            )
-            .await
-            .map_err(|e| {
-                Error::ExecutionErr(format!(
-                    "Failed to fetch S3 object for arg `{}`: {e}",
-                    arg.name
-                ))
-            })?;
-            snowflake_args.insert(arg.name.clone(), Value::String(json_text));
-        }
+                snowflake_args.insert(arg.name.clone(), Value::String(json_text));
+            }
+            Ok(())
+        };
+        crate::sql_s3_input::materialize_under_job_poller(
+            job,
+            conn,
+            mem_peak,
+            canceled_by,
+            worker_name,
+            occupancy_metrics,
+            materialize_f,
+        )
+        .await?;
     }
 
     let inline_db_res_path = parse_db_resource(&query);

backend/windmill-worker/src/sql_s3_input.rs

@@ -11,9 +11,67 @@
 //! executor binds the bare `s3://...` URI instead.
 
 use anyhow::Context;
+use std::future::Future;
 use windmill_common::client::AuthedClient;
+use windmill_common::error;
+use windmill_common::worker::Connection;
+use windmill_queue::{CanceledBy, MiniPulledJob};
 use windmill_types::s3::S3Object;
 
+use crate::common::OccupancyMetrics;
+use crate::handle_child::run_future_with_polling_update_job_poller;
+
+/// Run the `(s3object)` materialisation step under the job poller.
+///
+/// The poller is what refreshes `v2_job_runtime.ping`. Downloading and decoding the
+/// file can take minutes on large inputs, so running it unpolled leaves the ping
+/// stale and the zombie monitor restarts the job after `ZOMBIE_JOB_TIMEOUT`. The
+/// restart has no way to signal this phase, so the orphaned download keeps running
+/// concurrently with the retry.
+///
+/// This gives the phase its own `job.timeout` budget, on top of the one the query
+/// phase gets: a job can therefore occupy a worker for up to twice its timeout.
+/// That is deliberate — materialisation is otherwise unbounded — so don't collapse
+/// the two without giving this phase a budget of its own.
+pub(crate) async fn materialize_under_job_poller<T, Fut>(
+    job: &MiniPulledJob,
+    conn: &Connection,
+    mem_peak: &mut i32,
+    canceled_by: &mut Option<CanceledBy>,
+    worker_name: &str,
+    occupancy_metrics: &mut OccupancyMetrics,
+    fut: Fut,
+) -> error::Result<T>
+where
+    Fut: Future<Output = error::Result<T>>,
+{
+    run_future_with_polling_update_job_poller(
+        job.id,
+        job.timeout,
+        conn,
+        mem_peak,
+        canceled_by,
+        fut,
+        worker_name,
+        &job.workspace_id,
+        &mut Some(occupancy_metrics),
+        Box::pin(futures::stream::once(async { 0 })),
+    )
+    .await
+    .map_err(|e| match e {
+        // The poller labels its timeout arm as a query timeout, which would point
+        // the user at their SQL rather than at the input file. Relabelling on the
+        // message is best-effort: if it stops matching, the error is merely as
+        // unhelpful as it was before.
+        error::Error::ExecutionErr(m) if m.starts_with("Query timeout") => {
+            error::Error::ExecutionErr(format!(
+                "Timed out downloading and decoding the (s3object) argument: {m}"
+            ))
+        }
+        e => e,
+    })
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 enum InputFormat {
     Json,