From 045a52cdedcc27703d0b1983a324136b86a9eb1b Mon Sep 17 00:00:00 2001 From: yacine Bouraroui Date: Tue, 1 Oct 2024 19:47:07 +0200 Subject: [PATCH] query args from job_args table first --- backend/windmill-api/src/inputs.rs | 32 +++++------- backend/windmill-common/src/lib.rs | 1 + backend/windmill-common/src/macros.rs | 37 +++++++++++++ backend/windmill-worker/src/worker_flow.rs | 61 +++++++++------------- 4 files changed, 76 insertions(+), 55 deletions(-) create mode 100644 backend/windmill-common/src/macros.rs diff --git a/backend/windmill-api/src/inputs.rs b/backend/windmill-api/src/inputs.rs index e971056372..c1190e83b0 100644 --- a/backend/windmill-api/src/inputs.rs +++ b/backend/windmill-api/src/inputs.rs @@ -21,11 +21,7 @@ use std::{ vec, }; use windmill_common::{ - db::UserDB, - error::JsonResult, - jobs::JobKind, - scripts::to_i64, - utils::{not_found_if_none, paginate, Pagination}, + db::UserDB, error::JsonResult, jobs::JobKind, query_scalar_with_fallback, scripts::to_i64, utils::{not_found_if_none, paginate, Pagination} }; pub fn workspaced_service() -> Router { Router::new() @@ -178,6 +174,7 @@ struct GetArgs { input: Option, allow_large: Option, } + async fn get_args_from_history_or_saved_input( authed: ApiAuthed, Extension(user_db): Extension, @@ -185,6 +182,7 @@ async fn get_args_from_history_or_saved_input( Path((w_id, job_or_input_id)): Path<(String, Uuid)>, ) -> JsonResult> { let mut tx = user_db.begin(&authed).await?; + let result_o = if let Some(input) = g.input { if input { sqlx::query_scalar!( @@ -196,24 +194,20 @@ async fn get_args_from_history_or_saved_input( .fetch_optional(&mut *tx) .await? } else { - sqlx::query_scalar!( + query_scalar_with_fallback!(tx, + "SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM job_args WHERE id = $1 AND workspace_id = $2", "SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM completed_job WHERE id = $1 AND workspace_id = $2", - job_or_input_id, + job_or_input_id, w_id, - g.allow_large.unwrap_or(true) - ) - .fetch_optional(&mut *tx) - .await? + g.allow_large.unwrap_or(true))? } } else { - sqlx::query_scalar!( - "SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM completed_job WHERE id = $1 AND workspace_id = $2 UNION ALL SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM input WHERE id = $1 AND workspace_id = $2", - job_or_input_id, - w_id, - g.allow_large.unwrap_or(true) - ) - .fetch_optional(&mut *tx) - .await? + query_scalar_with_fallback!(tx, + "SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM job_args WHERE id = $1 AND workspace_id = $2 UNION ALL SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM input WHERE id = $1 AND workspace_id = $2", + "SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM completed_job WHERE id = $1 AND workspace_id = $2 UNION ALL SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM input WHERE id = $1 AND workspace_id = $2", + job_or_input_id, + w_id, + g.allow_large.unwrap_or(true))? }; tx.commit().await?; diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index cd8059251e..ce3b6a4e0b 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +pub mod macros; use std::{ net::SocketAddr, sync::{atomic::AtomicBool, Arc}, diff --git a/backend/windmill-common/src/macros.rs b/backend/windmill-common/src/macros.rs new file mode 100644 index 0000000000..b2f0263c3e --- /dev/null +++ b/backend/windmill-common/src/macros.rs @@ -0,0 +1,37 @@ +#[macro_export] +macro_rules! fetch_one_with_fallback { + ($db:expr, $query_method:ident, $fetch_method:ident ,$row_type:ty, $query:literal, $table:literal || $fallback_table:literal, $( $param:expr ),* ) => {{ + let primary_query = sqlx::$query_method::<_, $row_type>(const_format::formatcp!($query, $table)) + $(.bind($param))* + .$fetch_method($db) + .await; + + if let Err(sqlx::Error::RowNotFound) = primary_query { + tracing::info!("Data not found in job_params, falling back to fetching from $fallback_table"); + sqlx::$query_method::<_, $row_type>(const_format::formatcp!($query, $fallback_table)) + $(.bind($param))* + .$fetch_method($db) + .await + } else { + primary_query + } + }}; +} + +#[macro_export] +macro_rules! query_scalar_with_fallback { + ($tx:expr, $query:literal, $fallback_query:literal,$( $param:expr ),* ) => {{ + let primary_query = sqlx::query_scalar!($query,$($param),*) + .fetch_optional(&mut *$tx) + .await; + + if let Err(sqlx::Error::RowNotFound) = primary_query { + tracing::info!("Data not found in job_args, falling back to fetching from queue"); + sqlx::query_scalar!($fallback_query, $($param),*) + .fetch_optional(&mut *$tx) + .await + } else { + primary_query + } + }}; +} diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index dd5e50af2a..a9d51e132c 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -29,7 +29,6 @@ use sqlx::FromRow; use tokio::sync::mpsc::Sender; use tracing::instrument; use uuid::Uuid; -use windmill_common::add_time; use windmill_common::auth::JobPerms; #[cfg(feature = "benchmark")] use windmill_common::bench::BenchmarkIter; @@ -43,6 +42,7 @@ use windmill_common::jobs::{ RawCode, ENTRYPOINT_OVERRIDE, }; use windmill_common::worker::to_raw_value; +use windmill_common::{add_time, fetch_one_with_fallback}; use windmill_common::{ error::{self, to_anyhow, Error}, flow_status::{ @@ -61,32 +61,14 @@ type DB = sqlx::Pool; use windmill_queue::{canceled_job_to_result, get_queued_job_tx, push, QueueTransaction}; -macro_rules! fetch_one_with_fallback { - ($db:expr, $fetch_method:ident, $row_type:ty, $query:literal, $table:literal ,$( $param:expr ),* ) => {{ - let primary_query = sqlx::$fetch_method::<_, $row_type>(const_format::formatcp!($query, $table)) - $(.bind($param))* - .fetch_one($db) - .await; - - if let Err(sqlx::Error::RowNotFound) = primary_query { - tracing::info!("Data not found in job_params, falling back to fetching from queue"); - sqlx::$fetch_method::<_, $row_type>(const_format::formatcp!($query, "queue")) - $(.bind($param))* - .fetch_one($db) - .await - } else { - primary_query - } - }}; -} - async fn get_args_from_job_id(db: &DB, id: &Uuid) -> Result { let args = fetch_one_with_fallback!( db, query_as, + fetch_one, RowArgs, "SELECT args FROM {} WHERE id = $1", - "job_args", + "job_args" || "queue", id ); @@ -334,7 +316,12 @@ pub async fn update_flow_status_after_job_completion_internal< }; let is_flow = if let Some(step) = step { - fetch_one_with_fallback!(db, query_scalar, Option, "SELECT raw_flow->'modules'->($1)->'value'->>'type' = 'flow' FROM {} WHERE id = $2","job_params", + fetch_one_with_fallback!(db, + query_scalar, + fetch_one, + Option, + "SELECT raw_flow->'modules'->($1)->'value'->>'type' = 'flow' FROM {} WHERE id = $2", + "job_params" || "queue", step as i32, &flow).map_err(|e| { Error::InternalErr(format!("error during retrieval of step's type: {e:#}")) @@ -459,7 +446,7 @@ pub async fn update_flow_status_after_job_completion_internal< None }; - let nindex = if let Some(position) = position { + let nindex = if let Some(position) = position { sqlx::query_scalar!( "UPDATE queue SET flow_status = JSONB_SET( @@ -1283,9 +1270,10 @@ async fn has_failure_module<'c>(flow: Uuid, db: &DB) -> Result { fetch_one_with_fallback!( db, query_scalar, + fetch_one, Option, "SELECT raw_flow->'failure_module' != 'null'::jsonb FROM {} WHERE id = $1", - "job_params", + "job_params" || "queue", flow ) .map_err(|e| { @@ -2233,18 +2221,19 @@ async fn push_next_flow_job ); Ok(Marc::new(hm)) } else if let Some(id) = get_args_from_id { - let row = sqlx::query_as::<_, RawArgs>( - "SELECT args FROM completed_job WHERE id = $1 AND workspace_id = $2", - ) - .bind(id) - .bind(&flow_job.workspace_id) - .fetch_optional(db) - .await?; - + let row = fetch_one_with_fallback!( + db, + query_as, + fetch_optional, + RowArgs, + "SELECT args FROM {} WHERE id = $1 AND workspace_id = $2", + "job_args" || "completed_job", + id, + &flow_job.workspace_id + )?; + if let Some(raw_args) = row { - Ok(Marc::new( - raw_args.args.map(|x| x.0).unwrap_or_else(HashMap::new), - )) + Ok(Marc::new(raw_args.args.map(|x| x.0).unwrap_or_default())) } else { Ok(Marc::new(HashMap::new())) } @@ -2361,7 +2350,7 @@ async fn push_next_flow_job let mut tx: QueueTransaction<'_, R> = (rsmq.clone(), db.begin().await?).into(); let nargs = args.as_ref(); - for i in (0..len).into_iter() { + for i in 0..len { if i % 100 == 0 && i != 0 { tracing::info!(id = %flow_job.id, root_id = %job_root, "pushed (non-commited yet) first {i} subflows of {len}"); sqlx::query!(