From 7d605e88b0d8c600a0710f7735f021d1b77bdb91 Mon Sep 17 00:00:00 2001 From: yacine Bouraroui Date: Mon, 30 Sep 2024 18:25:29 +0200 Subject: [PATCH] read ( args, raw_code, raw_flow) from separate table - wip --- backend/windmill-api/src/jobs.rs | 31 +++---- backend/windmill-common/src/worker.rs | 4 +- backend/windmill-queue/src/jobs.rs | 32 +++++-- backend/windmill-worker/src/worker_flow.rs | 98 +++++++++++----------- 4 files changed, 90 insertions(+), 75 deletions(-) diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index 156be41b26..6113bdeaf7 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -11,7 +11,6 @@ use axum::http::HeaderValue; use quick_cache::sync::Cache; use serde_json::value::RawValue; use sqlx::Pool; -use windmill_common::error::JsonResult; use std::collections::HashMap; #[cfg(feature = "prometheus")] use std::sync::atomic::Ordering; @@ -19,6 +18,7 @@ use tokio::io::AsyncReadExt; #[cfg(feature = "prometheus")] use tokio::time::Instant; use tower::ServiceBuilder; +use windmill_common::error::JsonResult; use windmill_common::flow_status::{JobResult, RestartedFrom}; use windmill_common::jobs::{ format_completed_job_result, format_result, CompletedJobWithFormattedResult, FormattedResult, @@ -70,7 +70,10 @@ use windmill_common::{ oauth2::HmacSha256, scripts::{ScriptHash, ScriptLang}, users::username_to_permissioned_as, - utils::{not_found_if_none, now_from_db, paginate, paginate_without_limits, require_admin, Pagination, StripPath}, + utils::{ + not_found_if_none, now_from_db, paginate, paginate_without_limits, require_admin, + Pagination, StripPath, + }, }; #[cfg(all(feature = "enterprise", feature = "parquet"))] @@ -293,11 +296,8 @@ pub fn workspace_unauthed_service() -> Router { pub fn global_root_service() -> Router { Router::new() - .route("/db_clock", get(get_db_clock)) - .route( - "/completed/count_by_tag", - get(count_by_tag), - ) + .route("/db_clock", get(get_db_clock)) + .route("/completed/count_by_tag", get(count_by_tag)) } #[derive(Deserialize)] @@ -3234,7 +3234,7 @@ async fn run_wait_result( }; let fast_poll_duration = *WAIT_RESULT_FAST_POLL_DURATION_SECS as u64 * 1000; - let mut accumulated_delay = 0 as u64; + let mut accumulated_delay = 0_u64; loop { if let Some(node_id_for_empty_return) = node_id_for_empty_return.as_ref() { @@ -4344,11 +4344,8 @@ async fn add_batch_jobs( tx = PushIsolationLevel::Transaction(ntx); uuids.push(uuid); } - match tx { - PushIsolationLevel::Transaction(tx) => { - tx.commit().await?; - } - _ => (), + if let PushIsolationLevel::Transaction(tx) = tx { + tx.commit().await?; } return Ok(Json(uuids)); } @@ -4683,8 +4680,8 @@ async fn get_job_update( .fetch_optional(&db) .await?; - let progress: Option = if get_progress == Some(true){ - sqlx::query_scalar!( + let progress: Option = if get_progress == Some(true) { + sqlx::query_scalar!( "SELECT scalar_int FROM job_stats WHERE workspace_id = $1 AND job_id = $2 AND metric_id = $3", &w_id, job_id, @@ -5115,8 +5112,6 @@ async fn get_completed_job_result( Ok(Json(result).into_response()) } - - #[derive(Deserialize)] struct CountByTagQuery { horizon_secs: Option, @@ -5130,7 +5125,7 @@ struct TagCount { } async fn count_by_tag( - ApiAuthed { email, ..}: ApiAuthed, + ApiAuthed { email, .. }: ApiAuthed, Extension(db): Extension, Query(query): Query, ) -> JsonResult> { diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index 9464eccdcc..1bfca16256 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -90,7 +90,7 @@ lazy_static::lazy_static! { } pub async fn make_suspended_pull_query(wc: &WorkerConfig) { - if wc.worker_tags.len() == 0 { + if wc.worker_tags.is_empty() { tracing::error!("Empty tags in worker tags, skipping"); return; } @@ -122,7 +122,7 @@ pub async fn make_suspended_pull_query(wc: &WorkerConfig) { pub async fn make_pull_query(wc: &WorkerConfig) { let mut queries = vec![]; for tags in wc.priority_tags_sorted.iter() { - if tags.tags.len() == 0 { + if tags.tags.is_empty() { tracing::error!("Empty tags in priority tags, skipping"); continue; } diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index eca874aa63..74383b6b82 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -2291,8 +2291,7 @@ pub async fn get_result_by_id( Some(job) => { let restarted_from = windmill_common::utils::not_found_if_none( job.parse_flow_status() - .map(|status| status.restarted_from) - .flatten(), + .and_then(|status| status.restarted_from), "Id not found in the result's mapping of the root job and root job had no restarted from information", format!("parent: {}, root: {}, id: {}", flow_id, job.id, node_id), )?; @@ -2351,8 +2350,7 @@ pub async fn get_result_by_id_from_running_flow( let job_result = flow_job_result .leaf_jobs - .map(|x| serde_json::from_str(x.get()).ok()) - .flatten(); + .and_then(|x| serde_json::from_str(x.get()).ok()); if job_result.is_none() && flow_job_result.parent_job.is_some() { let parent_job = flow_job_result.parent_job.unwrap(); @@ -2642,7 +2640,7 @@ pub struct PushArgsOwned { pub args: HashMap>, } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct PushArgs<'c> { pub extra: Option>>, pub args: &'c HashMap>, @@ -3893,6 +3891,10 @@ pub async fn push<'c, 'd, R: rsmq_async::RsmqConnection + Send + 'c>( }; tracing::debug!("Pushing job {job_id} with tag {tag}, schedule_path {schedule_path:?}, script_path: {script_path:?}, email {email}, workspace_id {workspace_id}"); + + let raw_flow = raw_flow.map(Json); + let args = Json(args); + let uuid = sqlx::query_scalar!( "INSERT INTO queue (workspace_id, id, running, parent_job, created_by, permissioned_as, scheduled_for, @@ -3911,12 +3913,12 @@ pub async fn push<'c, 'd, R: rsmq_async::RsmqConnection + Send + 'c>( scheduled_for_o, script_hash, script_path.clone(), - raw_code, + raw_code.clone(), raw_lock, - Json(args) as Json, + args.clone() as Json, job_kind.clone() as JobKind, schedule_path, - raw_flow.map(Json) as Option>, + raw_flow.clone() as Option>, flow_status.map(Json) as Option>, is_flow_step, language as Option, @@ -3937,6 +3939,20 @@ pub async fn push<'c, 'd, R: rsmq_async::RsmqConnection + Send + 'c>( .await .map_err(|e| Error::InternalErr(format!("Could not insert into queue {job_id} with tag {tag}, schedule_path {schedule_path:?}, script_path: {script_path:?}, email {email}, workspace_id {workspace_id}: {e:#}")))?; + // insert into args queue + sqlx::query!( + r#" + INSERT INTO job_params (id, args, raw_code, raw_flow) + VALUES ($1, $2, $3, $4) + "#, + uuid, + args as Json, + raw_code, + raw_flow as Option> + ) + .execute(&mut tx) + .await?; + tracing::debug!("Pushed {job_id}"); // TODO: technically the job isn't queued yet, as the transaction can be rolled back. Should be solved when moving these metrics to the queue abstraction. #[cfg(feature = "prometheus")] diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 6474af6d46..cf979fc340 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -61,6 +61,37 @@ 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, $( $param:expr ),* ) => {{ + let primary_query = sqlx::$fetch_method::<_, $row_type>(const_format::formatcp!($query, "job_params")) + $(.bind($param))* + .fetch_one($db) + .await; + + if let Err(sqlx::Error::RowNotFound) = primary_query { + // Fallback to the second query if the first query returns no result + 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, + RowArgs, + "SELECT args FROM {} WHERE id = $1", + id + ); + + args.map_err(|e| Error::InternalErr(format!("retrieval of args from state: {e:#}"))) +} + // #[instrument(level = "trace", skip_all)] pub async fn update_flow_status_after_job_completion< R: rsmq_async::RsmqConnection + Send + Sync + Clone, @@ -229,10 +260,10 @@ pub async fn update_flow_status_after_job_completion_internal< })?; let old_status = serde_json::from_str::(old_status_json.flow_status.get()) - .or_else(|e| { - Err(Error::InternalErr(format!( + .map_err(|e| { + Error::InternalErr(format!( "requiring status to be parsable as FlowStatus: {e:?}" - ))) + )) })?; let current_module = if let Some(x) = old_status_json.current_module { @@ -289,7 +320,7 @@ pub async fn update_flow_status_after_job_completion_internal< // 0 length flows are not failure steps let is_failure_step = - old_status.step >= old_status.modules.len() as i32 && old_status.modules.len() > 0; + old_status.step >= old_status.modules.len() as i32 && !old_status.modules.is_empty(); let (mut stop_early, mut skip_if_stop_early, continue_on_error) = if let Some(se) = stop_early_override @@ -302,14 +333,9 @@ pub async fn update_flow_status_after_job_completion_internal< }; let is_flow = if let Some(step) = step { - sqlx::query_scalar!( - "SELECT raw_flow->'modules'->($1)->'value'->>'type' = 'flow' FROM queue WHERE id = $2", + fetch_one_with_fallback!(db, query_scalar, Option, "SELECT raw_flow->'modules'->($1)->'value'->>'type' = 'flow' FROM {} WHERE id = $2", step as i32, - &flow - ) - .fetch_one(db) - .await - .map_err(|e| { + &flow).map_err(|e| { Error::InternalErr(format!("error during retrieval of step's type: {e:#}")) })? .unwrap_or(false) @@ -342,19 +368,9 @@ pub async fn update_flow_status_after_job_completion_internal< } _ => None, }; - let args = sqlx::query_as::<_, RowArgs>( - "SELECT - args - FROM queue - WHERE id = $2", - ) - .bind(old_status.step) - .bind(flow) - .fetch_one(db) - .await - .map_err(|e| { - Error::InternalErr(format!("retrieval of args from state: {e:#}")) - })?; + + let args = get_args_from_job_id(db, &flow).await?; + compute_bool_from_expr( expr.to_string(), Marc::new(args.args.unwrap_or_default().0), @@ -489,7 +505,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( @@ -657,7 +673,7 @@ pub async fn update_flow_status_after_job_completion_internal< flow_jobs_success, flow_jobs, .. - } if branch.to_owned() < len - 1 && (success || skip_branch_failure) => { + } if *branch < len - 1 && (success || skip_branch_failure) => { if let Some(jobs) = flow_jobs { set_success_in_flow_job_success( flow_jobs_success, @@ -854,19 +870,7 @@ pub async fn update_flow_status_after_job_completion_internal< .as_ref() .and_then(|m| m.stop_after_all_iters_if.as_ref().map(|x| x.expr.clone())) { - let args = sqlx::query_as::<_, RowArgs>( - "SELECT - args - FROM queue - WHERE id = $2", - ) - .bind(old_status.step) - .bind(flow) - .fetch_one(db) - .await - .map_err(|e| { - Error::InternalErr(format!("retrieval of args from state: {e:#}")) - })?; + let args = get_args_from_job_id(db, &flow).await?; let should_stop = compute_bool_from_expr( expr.to_string(), @@ -1275,14 +1279,13 @@ async fn compute_skip_branchall_failure<'c>( } async fn has_failure_module<'c>(flow: Uuid, db: &DB) -> Result { - sqlx::query_scalar::<_, Option>( - "SELECT raw_flow->'failure_module' != 'null'::jsonb - FROM queue - WHERE id = $1", + fetch_one_with_fallback!( + db, + query_scalar, + Option, + "SELECT raw_flow->'failure_module' != 'null'::jsonb FROM {} WHERE id = $1", + flow ) - .bind(flow) - .fetch_one(db) - .await .map_err(|e| { Error::InternalErr(format!( "error during retrieval of has_failure_module: {e:#}" @@ -1622,6 +1625,7 @@ fn potentially_crash_for_testing() { lazy_static::lazy_static! { pub static ref EHM: HashMap> = HashMap::new(); } + // #[async_recursion] // #[instrument(level = "trace", skip_all)] async fn push_next_flow_job( @@ -1709,7 +1713,7 @@ async fn push_next_flow_job flow_job.workspace_id.as_str(), flow_job.id ).fetch_all(db).await?; - if overlapping.len() > 0 { + if !overlapping.is_empty() { let overlapping_str = overlapping .iter() .map(|x| x.to_string())