refactor future wrapper for queries

backend/ee-repo-ref.txt

@@ -1 +1 @@
-3de554fc89db8f9c6cc544c6a2ba004e65a7c133
+fd0c4dab2ae54bed1c7cd7041669fc3463cd163e

backend/windmill-worker/src/common.rs

@@ -14,6 +14,7 @@ use serde_json::{json, Value};
 use sqlx::{Pool, Postgres};
 use tokio::process::Command;
 use tokio::{fs::File, io::AsyncReadExt};
+use windmill_common::error::to_anyhow;
 use windmill_common::s3_helpers::{
     get_etag_or_empty, AzureBlobResource, LargeFileStorage, ObjectStoreResource, S3Object,
     S3Resource,
@@ -447,6 +448,53 @@ async fn get_mem_peak(pid: Option<u32>, nsjail: bool) -> i32 {
     }
 }
 
+pub async fn run_future_with_polling_update_job_poller<Fut, T>(
+    job_id: Uuid,
+    timeout: Option<i32>,
+    db: &DB,
+    mem_peak: &mut i32,
+    canceled_by_ref: &mut Option<CanceledBy>,
+    result_f: Fut,
+    worker_name: &str,
+    w_id: &str,
+) -> anyhow::Result<T>
+where
+    Fut: Future<Output = anyhow::Result<T>>,
+{
+    let (tx, rx) = broadcast::channel::<()>(3);
+
+    let update_job = update_job_poller(
+        job_id,
+        db,
+        mem_peak,
+        canceled_by_ref,
+        || async { 0 },
+        worker_name,
+        w_id,
+        rx,
+    );
+
+    let timeout_ms = u64::try_from(
+        resolve_job_timeout(&db, &w_id, job_id, timeout)
+            .await
+            .0
+            .as_millis(),
+    )
+    .unwrap_or(200000);
+
+    let rows = tokio::select! {
+        biased;
+        result = tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), result_f) => result
+        .map_err(|e| {
+            tracing::error!("Query timeout: {}", e);
+            Error::ExecutionErr("Query timeout".to_string())
+        })?,
+        _ = update_job, if job_id != Uuid::nil() => Err(Error::ExecutionErr("Job cancelled".to_string())).map_err(to_anyhow)?,
+    }?;
+    drop(tx);
+    Ok(rows)
+}
+
 pub async fn update_job_poller<F, Fut>(
     job_id: Uuid,
     db: &DB,

backend/windmill-worker/src/mysql_executor.rs

@@ -12,8 +12,12 @@ use windmill_common::{
     jobs::QueuedJob,
 };
 use windmill_parser_sql::{parse_mysql_sig, RE_ARG_MYSQL_NAMED};
+use windmill_queue::CanceledBy;
 
-use crate::{common::build_args_map, AuthedClientBackgroundTask};
+use crate::{
+    common::{build_args_map, run_future_with_polling_update_job_poller},
+    AuthedClientBackgroundTask,
+};
 
 #[derive(Deserialize)]
 struct MysqlDatabase {
@@ -30,6 +34,9 @@ pub async fn do_mysql(
     client: &AuthedClientBackgroundTask,
     query: &str,
     db: &sqlx::Pool<sqlx::Postgres>,
+    mem_peak: &mut i32,
+    canceled_by: &mut Option<CanceledBy>,
+    worker_name: &str,
 ) -> windmill_common::error::Result<Box<RawValue>> {
     let args = build_args_map(job, client, db).await?.map(Json);
     let job_args = if args.is_some() {
@@ -129,28 +136,47 @@ pub async fn do_mysql(
             _ => {}
         }
     }
-    let rows: Vec<Row> = conn
-        .exec(
-            query,
-            match statement_values {
-                Params::Positional(v) => Params::Positional(v),
-                Params::Named(m) => Params::Named(m),
-                _ => Params::Empty,
-            },
-        )
-        .await
-        .map_err(to_anyhow)?;
-    let rows = rows
-        .into_iter()
-        .map(|x| convert_row_to_value(x))
-        .collect::<Vec<serde_json::Value>>();
+
+    let result_f = async {
+        let rows: Vec<Row> = conn
+            .exec(
+                query,
+                match statement_values {
+                    Params::Positional(v) => Params::Positional(v),
+                    Params::Named(m) => Params::Named(m),
+                    _ => Params::Empty,
+                },
+            )
+            .await
+            .map_err(to_anyhow)?;
+        Ok(rows
+            .into_iter()
+            .map(|x| convert_row_to_value(x))
+            .collect::<Vec<serde_json::Value>>())
+            as Result<Vec<serde_json::Value>, anyhow::Error>
+    };
+
+    let result = run_future_with_polling_update_job_poller(
+        job.id,
+        job.timeout,
+        db,
+        mem_peak,
+        canceled_by,
+        result_f,
+        worker_name,
+        &job.workspace_id,
+    )
+    .await?;
 
     drop(conn);
 
     pool.disconnect().await.map_err(to_anyhow)?;
 
+    let raw_result = windmill_common::worker::to_raw_value(&json!(result));
+    *mem_peak = (raw_result.get().len() / 1000) as i32;
+
     // And then check that we got back the same string we sent over.
-    return Ok(windmill_common::worker::to_raw_value(&json!(rows)));
+    return Ok(raw_result);
 }
 
 fn string_date_to_mysql_date(s: &str) -> mysql_async::Value {

backend/windmill-worker/src/pg_executor.rs

@@ -13,7 +13,7 @@ use serde::Deserialize;
 use serde_json::value::RawValue;
 use serde_json::Map;
 use serde_json::Value;
-use tokio::sync::{broadcast, Mutex};
+use tokio::sync::Mutex;
 use tokio_postgres::types::IsNull;
 use tokio_postgres::{
     types::{to_sql_checked, ToSql},
@@ -31,7 +31,7 @@ use windmill_parser::Typ;
 use windmill_parser_sql::parse_pgsql_sig;
 use windmill_queue::CanceledBy;
 
-use crate::common::{build_args_values, resolve_job_timeout, update_job_poller};
+use crate::common::{build_args_values, run_future_with_polling_update_job_poller};
 use crate::AuthedClientBackgroundTask;
 use bytes::{Buf, BytesMut};
 use lazy_static::lazy_static;
@@ -206,37 +206,17 @@ pub async fn do_postgresql(
             .collect::<Result<Vec<_>, _>>()?) as anyhow::Result<Vec<serde_json::Value>>
     };
 
-    let (tx, rx) = broadcast::channel::<()>(3);
-
-    let update_job = update_job_poller(
+    let result = run_future_with_polling_update_job_poller(
         job.id,
+        job.timeout,
         db,
         mem_peak,
         canceled_by,
-        || async { 0 },
+        result_f,
         worker_name,
         &job.workspace_id,
-        rx,
-    );
-
-    let timeout_ms = u64::try_from(
-        resolve_job_timeout(&db, &job.workspace_id, job.id, job.timeout)
-            .await
-            .0
-            .as_millis(),
     )
-    .unwrap_or(200000);
-    tracing::info!("Query timeout: {}ms", timeout_ms);
-    let result = tokio::select! {
-        biased;
-        result = tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), result_f) => result
-        .map_err(|e| {
-            tracing::error!("Query timeout: {}", e);
-            Error::ExecutionErr("Query timeout".to_string())
-        })?,
-        _ = update_job, if job.id != Uuid::nil() => Err(Error::ExecutionErr("Job cancelled".to_string())).map_err(to_anyhow)?,
-    }?;
-    drop(tx);
+    .await?;
 
     RUNNING.store(false, std::sync::atomic::Ordering::Relaxed);
 

backend/windmill-worker/src/worker.rs

@@ -2893,7 +2893,16 @@ async fn handle_code_execution_job(
         )
         .await;
     } else if language == Some(ScriptLang::Mysql) {
-        return do_mysql(job, &client, &inner_content, db).await;
+        return do_mysql(
+            job,
+            &client,
+            &inner_content,
+            db,
+            mem_peak,
+            canceled_by,
+            worker_name,
+        )
+        .await;
     } else if language == Some(ScriptLang::Bigquery) {
         #[cfg(not(feature = "enterprise"))]
         {