From 2ecf44f44fb0f7e3c236eab3762780fc7e334209 Mon Sep 17 00:00:00 2001 From: HugoCasa Date: Wed, 3 Jul 2024 21:51:59 +0200 Subject: [PATCH] fix(backend): switch job run to user db (#4017) --- backend/windmill-api/src/jobs.rs | 85 +++++++++++++++------- backend/windmill-common/src/jobs.rs | 10 +-- backend/windmill-common/src/lib.rs | 4 +- backend/windmill-worker/src/worker_flow.rs | 2 +- 4 files changed, 68 insertions(+), 33 deletions(-) diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index b0bf34c104..a3298b8e2d 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -77,7 +77,7 @@ use windmill_common::{METRICS_DEBUG_ENABLED, METRICS_ENABLED}; use windmill_common::{get_latest_deployed_hash_for_path, BASE_URL}; use windmill_queue::{ cancel_job, get_queued_job, get_result_by_id_from_running_flow, job_is_complete, push, - DecodeQueries, PushArgs, PushIsolationLevel, + DecodeQueries, PushArgs, PushIsolationLevel, QueueTransaction, }; #[cfg(feature = "prometheus")] @@ -529,8 +529,8 @@ pub async fn get_path_for_hash<'c>( Ok(path) } -pub async fn get_path_tag_limits_cache_for_hash( - db: &DB, +pub async fn get_path_tag_limits_cache_for_hash<'c, R: rsmq_async::RsmqConnection + Send>( + tx: &mut QueueTransaction<'c, R>, w_id: &str, hash: i64, ) -> error::Result<( @@ -551,13 +551,15 @@ pub async fn get_path_tag_limits_cache_for_hash( hash, w_id ) - .fetch_one(db) + .fetch_optional(tx) .await .map_err(|e| { Error::InternalErr(format!( "querying getting path for hash {hash} in {w_id}: {e:#}" )) - })?; + })?.ok_or_else(|| Error::NotFound(format!( + "deployed script not found at hash {hash} in workspace {w_id}" + )))?; Ok(( script.path, script.tag, @@ -2641,21 +2643,27 @@ pub async fn run_flow_by_path( let flow_path = flow_path.to_path(); check_scopes(&authed, || format!("run:flow/{flow_path}"))?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let (tag, dedicated_worker) = sqlx::query!( "SELECT tag, dedicated_worker from flow WHERE path = $1 and workspace_id = $2", flow_path, w_id ) - .fetch_optional(&db) + .fetch_optional(&mut tx) .await? .map(|x| (x.tag, x.dedicated_worker)) - .unwrap_or_else(|| (None, None)); + .ok_or_else(|| { + Error::NotFound(format!( + "flow not found at path {flow_path} in workspace {w_id}" + )) + })?; let tag = run_query.tag.clone().or(tag); check_tag_available_for_workspace(&w_id, &tag).await?; let scheduled_for = run_query.get_scheduled_for(&db).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, tx, @@ -2720,12 +2728,14 @@ pub async fn restart_flow( ) -> error::Result<(StatusCode, String)> { check_license_key_valid().await?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let completed_job = sqlx::query_as::<_, CompletedJob>( "SELECT *, result->'wm_labels' as labels from completed_job WHERE id = $1 and workspace_id = $2", ) .bind(job_id) .bind(&w_id) - .fetch_optional(&db) + .fetch_optional(&mut tx) .await? .with_context(|| "Unable to find completed job with the given job UUID")?; @@ -2740,7 +2750,8 @@ pub async fn restart_flow( .unwrap_or_else(PushArgs::empty); let scheduled_for = run_query.get_scheduled_for(&db).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, @@ -2791,13 +2802,16 @@ pub async fn run_script_by_path( check_scopes(&authed, || format!("run:script/{script_path}"))?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let (job_payload, tag, _delete_after_use, timeout) = - script_path_to_payload(script_path, &db, &w_id).await?; + script_path_to_payload(script_path, &mut tx, &w_id).await?; let scheduled_for = run_query.get_scheduled_for(&db).await?; let tag = run_query.tag.clone().or(tag); check_tag_available_for_workspace(&w_id, &tag).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, @@ -2841,6 +2855,8 @@ pub async fn run_workflow_as_code( check_license_key_valid().await?; check_tag_available_for_workspace(&w_id, &run_query.tag).await?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let job = get_queued_job(&job_id, &w_id, &db).await?; let job = not_found_if_none(job, "Queued Job", &job_id.to_string())?; let (job_payload, tag, _delete_after_use, timeout) = match job.job_kind { @@ -2863,7 +2879,7 @@ pub async fn run_workflow_as_code( None, run_query.timeout, ), - JobKind::Script => script_path_to_payload(job.script_path(), &db, &w_id).await?, + JobKind::Script => script_path_to_payload(job.script_path(), &mut tx, &w_id).await?, _ => return Err(anyhow::anyhow!("Not supported").into()), }; @@ -2874,7 +2890,8 @@ pub async fn run_workflow_as_code( let scheduled_for = run_query.get_scheduled_for(&db).await?; let tag = run_query.tag.clone().or(tag).or(Some(job.tag)); - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let tx = PushIsolationLevel::Transaction(tx); let (uuid, mut tx) = push( &db, @@ -3192,12 +3209,15 @@ pub async fn run_wait_result_job_by_path_get( let script_path = script_path.to_path(); check_scopes(&authed, || format!("run:script/{script_path}"))?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let (job_payload, tag, delete_after_use, timeout) = - script_path_to_payload(script_path, &db, &w_id).await?; + script_path_to_payload(script_path, &mut tx, &w_id).await?; let tag = run_query.tag.clone().or(tag); check_tag_available_for_workspace(&w_id, &tag).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, @@ -3312,12 +3332,15 @@ async fn run_wait_result_script_by_path_internal( let script_path = script_path.to_path(); check_scopes(&authed, || format!("run:script/{script_path}"))?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let (job_payload, tag, delete_after_use, timeout) = - script_path_to_payload(script_path, &db, &w_id).await?; + script_path_to_payload(script_path, &mut tx, &w_id).await?; let tag = run_query.tag.clone().or(tag); check_tag_available_for_workspace(&w_id, &tag).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, @@ -3367,6 +3390,8 @@ pub async fn run_wait_result_script_by_hash( check_queue_too_long(&db, run_query.queue_limit).await?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let hash = script_hash.0; let ( path, @@ -3380,7 +3405,7 @@ pub async fn run_wait_result_script_by_hash( priority, delete_after_use, timeout, - ) = get_path_tag_limits_cache_for_hash(&db, &w_id, hash).await?; + ) = get_path_tag_limits_cache_for_hash(&mut tx, &w_id, hash).await?; if let Some(run_query_cache_ttl) = run_query.cache_ttl { cache_ttl = Some(run_query_cache_ttl); } @@ -3388,7 +3413,8 @@ pub async fn run_wait_result_script_by_hash( let tag = run_query.tag.clone().or(tag); check_tag_available_for_workspace(&w_id, &tag).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, @@ -3467,6 +3493,8 @@ async fn run_wait_result_flow_by_path_internal( let flow_path = flow_path.to_path(); check_scopes(&authed, || format!("run:flow/{flow_path}"))?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let scheduled_for = run_query.get_scheduled_for(&db).await?; let (tag, dedicated_worker, early_return) = sqlx::query!( @@ -3478,14 +3506,19 @@ async fn run_wait_result_flow_by_path_internal( flow_path, w_id ) - .fetch_optional(&db) + .fetch_optional(&mut tx) .await? .map(|x| (x.tag, x.dedicated_worker, x.early_return)) - .unwrap_or_else(|| (None, None, None)); + .ok_or_else( + || Error::NotFound( + format!("flow not found at path {flow_path} in workspace {w_id}") + ) + )?; let tag = run_query.tag.clone().or(tag); check_tag_available_for_workspace(&w_id, &tag).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, @@ -4111,6 +4144,8 @@ pub async fn run_job_by_hash( #[cfg(feature = "enterprise")] check_license_key_valid().await?; + let mut tx: QueueTransaction<'_, _> = (rsmq, user_db.begin(&authed).await?).into(); + let hash = script_hash.0; let ( path, @@ -4124,7 +4159,7 @@ pub async fn run_job_by_hash( priority, _delete_after_use, // not taken into account in async endpoints timeout, - ) = get_path_tag_limits_cache_for_hash(&db, &w_id, hash).await?; + ) = get_path_tag_limits_cache_for_hash(&mut tx, &w_id, hash).await?; check_scopes(&authed, || format!("run:script/{path}"))?; if let Some(run_query_cache_ttl) = run_query.cache_ttl { cache_ttl = Some(run_query_cache_ttl); @@ -4133,7 +4168,7 @@ pub async fn run_job_by_hash( let tag = run_query.tag.clone().or(tag); check_tag_available_for_workspace(&w_id, &tag).await?; - let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + let tx = PushIsolationLevel::Transaction(tx); let (uuid, tx) = push( &db, diff --git a/backend/windmill-common/src/jobs.rs b/backend/windmill-common/src/jobs.rs index 300d204415..2590087496 100644 --- a/backend/windmill-common/src/jobs.rs +++ b/backend/windmill-common/src/jobs.rs @@ -384,9 +384,9 @@ type Tag = String; pub type DB = Pool; -pub async fn script_path_to_payload( +pub async fn script_path_to_payload<'e, E: sqlx::Executor<'e, Database = Postgres>>( script_path: &str, - db: &DB, + db: E, w_id: &str, ) -> error::Result<(JobPayload, Option, Option, Option)> { let (job_payload, tag, delete_after_use, script_timeout) = if script_path.starts_with("hub/") { @@ -472,13 +472,13 @@ pub async fn script_hash_to_tag_and_limits<'c>( )) } -pub async fn get_payload_tag_from_prefixed_path( +pub async fn get_payload_tag_from_prefixed_path<'e, E: sqlx::Executor<'e, Database = Postgres>>( path: &str, - db: &DB, + db: E, w_id: &str, ) -> Result<(JobPayload, Option), Error> { let (payload, tag, _, _) = if path.starts_with("script/") { - script_path_to_payload(path.strip_prefix("script/").unwrap(), &db, w_id).await? + script_path_to_payload(path.strip_prefix("script/").unwrap(), db, w_id).await? } else if path.starts_with("flow/") { let path = path.strip_prefix("flow/").unwrap().to_string(); let r = sqlx::query!( diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index eb2cf6eb6b..d4cd97409c 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -243,8 +243,8 @@ type Tag = String; pub type DB = Pool; -pub async fn get_latest_deployed_hash_for_path( - db: &DB, +pub async fn get_latest_deployed_hash_for_path<'e, E: sqlx::Executor<'e, Database = Postgres>>( + db: E, w_id: &str, script_path: &str, ) -> error::Result<( diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index ba8d5ada54..8b558bd5bb 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -3248,7 +3248,7 @@ async fn script_to_payload( tag_override: &Option, ) -> Result { let (payload, tag, delete_after_use, script_timeout) = if script_hash.is_none() { - script_path_to_payload(script_path, &db, &flow_job.workspace_id).await? + script_path_to_payload(script_path, db, &flow_job.workspace_id).await? } else { let hash = script_hash.clone().unwrap(); let mut tx: sqlx::Transaction<'_, sqlx::Postgres> = db.begin().await?;