From 09c60442baaa68b3f00d6129a216bd166e69bae1 Mon Sep 17 00:00:00 2001 From: Abel Lucas Date: Thu, 16 Jan 2025 20:16:24 +0100 Subject: [PATCH] backend: rework old `job` table usages --- ...fdaa1cbd3fe93b590b97c549ff2ad0016da1a.json | 66 ++++++ ...c6c6927ac0a6f6c26f11db69fbc4aa839cf64.json | 91 ++++++++ ...31981cb91ccabdaa786ebadb6d80e8a8b38db.json | 22 ++ ...7ac50acac152fc0a5315e90fce6d0abeac652.json | 67 ++++++ backend/windmill-api/src/jobs.rs | 43 ++-- backend/windmill-api/src/slack_approvals.rs | 17 +- backend/windmill-common/src/cache.rs | 134 +++-------- backend/windmill-queue/src/jobs.rs | 214 +++++++----------- backend/windmill-worker/src/worker.rs | 30 +-- backend/windmill-worker/src/worker_flow.rs | 13 +- 10 files changed, 387 insertions(+), 310 deletions(-) create mode 100644 backend/.sqlx/query-111ee00fa6661fd4d3bb1d2e567fdaa1cbd3fe93b590b97c549ff2ad0016da1a.json create mode 100644 backend/.sqlx/query-3af91bb37638473d3d97a3a6f7fc6c6927ac0a6f6c26f11db69fbc4aa839cf64.json create mode 100644 backend/.sqlx/query-4a4971ab285cc5ca0dc4026d75a31981cb91ccabdaa786ebadb6d80e8a8b38db.json create mode 100644 backend/.sqlx/query-ebd8349ebda8ecb57e15257233a7ac50acac152fc0a5315e90fce6d0abeac652.json diff --git a/backend/.sqlx/query-111ee00fa6661fd4d3bb1d2e567fdaa1cbd3fe93b590b97c549ff2ad0016da1a.json b/backend/.sqlx/query-111ee00fa6661fd4d3bb1d2e567fdaa1cbd3fe93b590b97c549ff2ad0016da1a.json new file mode 100644 index 0000000000..bc35c5b2f7 --- /dev/null +++ b/backend/.sqlx/query-111ee00fa6661fd4d3bb1d2e567fdaa1cbd3fe93b590b97c549ff2ad0016da1a.json @@ -0,0 +1,66 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT kind AS \"kind: JobKind\", runnable_id, runnable_path,\n flow_status AS \"flow_status: Json>\"\n FROM v2_job JOIN v2_job_completed USING (id)\n WHERE id = $1", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "kind: JobKind", + "type_info": { + "Custom": { + "name": "job_kind", + "kind": { + "Enum": [ + "script", + "preview", + "flow", + "dependencies", + "flowpreview", + "script_hub", + "identity", + "flowdependencies", + "http", + "graphql", + "postgresql", + "noop", + "appdependencies", + "deploymentcallback", + "singlescriptflow", + "flowscript", + "flownode", + "appscript" + ] + } + } + } + }, + { + "ordinal": 1, + "name": "runnable_id", + "type_info": "Int8" + }, + { + "ordinal": 2, + "name": "runnable_path", + "type_info": "Varchar" + }, + { + "ordinal": 3, + "name": "flow_status: Json>", + "type_info": "Jsonb" + } + ], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [ + false, + true, + true, + true + ] + }, + "hash": "111ee00fa6661fd4d3bb1d2e567fdaa1cbd3fe93b590b97c549ff2ad0016da1a" +} diff --git a/backend/.sqlx/query-3af91bb37638473d3d97a3a6f7fc6c6927ac0a6f6c26f11db69fbc4aa839cf64.json b/backend/.sqlx/query-3af91bb37638473d3d97a3a6f7fc6c6927ac0a6f6c26f11db69fbc4aa839cf64.json new file mode 100644 index 0000000000..1809e60091 --- /dev/null +++ b/backend/.sqlx/query-3af91bb37638473d3d97a3a6f7fc6c6927ac0a6f6c26f11db69fbc4aa839cf64.json @@ -0,0 +1,91 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT\n v2_queue.job_kind AS \"job_kind!: JobKind\",\n v2_queue.script_hash,\n v2_queue.raw_flow AS \"raw_flow: sqlx::types::Json>\",\n v2_completed_job.parent_job AS \"parent_job: Uuid\",\n v2_completed_job.created_at AS \"created_at!: chrono::NaiveDateTime\",\n v2_completed_job.created_by AS \"created_by!\",\n v2_queue.script_path,\n v2_queue.args AS \"args: sqlx::types::Json>\"\n FROM v2_queue\n JOIN v2_completed_job ON v2_completed_job.parent_job = v2_queue.id\n WHERE v2_completed_job.id = $1 AND v2_completed_job.workspace_id = $2\n LIMIT 1", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_kind!: JobKind", + "type_info": { + "Custom": { + "name": "job_kind", + "kind": { + "Enum": [ + "script", + "preview", + "flow", + "dependencies", + "flowpreview", + "script_hub", + "identity", + "flowdependencies", + "http", + "graphql", + "postgresql", + "noop", + "appdependencies", + "deploymentcallback", + "singlescriptflow", + "flowscript", + "flownode", + "appscript" + ] + } + } + } + }, + { + "ordinal": 1, + "name": "script_hash", + "type_info": "Int8" + }, + { + "ordinal": 2, + "name": "raw_flow: sqlx::types::Json>", + "type_info": "Jsonb" + }, + { + "ordinal": 3, + "name": "parent_job: Uuid", + "type_info": "Uuid" + }, + { + "ordinal": 4, + "name": "created_at!: chrono::NaiveDateTime", + "type_info": "Timestamptz" + }, + { + "ordinal": 5, + "name": "created_by!", + "type_info": "Varchar" + }, + { + "ordinal": 6, + "name": "script_path", + "type_info": "Varchar" + }, + { + "ordinal": 7, + "name": "args: sqlx::types::Json>", + "type_info": "Jsonb" + } + ], + "parameters": { + "Left": [ + "Uuid", + "Text" + ] + }, + "nullable": [ + true, + true, + true, + true, + true, + true, + true, + true + ] + }, + "hash": "3af91bb37638473d3d97a3a6f7fc6c6927ac0a6f6c26f11db69fbc4aa839cf64" +} diff --git a/backend/.sqlx/query-4a4971ab285cc5ca0dc4026d75a31981cb91ccabdaa786ebadb6d80e8a8b38db.json b/backend/.sqlx/query-4a4971ab285cc5ca0dc4026d75a31981cb91ccabdaa786ebadb6d80e8a8b38db.json new file mode 100644 index 0000000000..473149f877 --- /dev/null +++ b/backend/.sqlx/query-4a4971ab285cc5ca0dc4026d75a31981cb91ccabdaa786ebadb6d80e8a8b38db.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT flow_status AS \"flow_status: Json>\"\n FROM v2_job_completed WHERE id = $1", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "flow_status: Json>", + "type_info": "Jsonb" + } + ], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [ + true + ] + }, + "hash": "4a4971ab285cc5ca0dc4026d75a31981cb91ccabdaa786ebadb6d80e8a8b38db" +} diff --git a/backend/.sqlx/query-ebd8349ebda8ecb57e15257233a7ac50acac152fc0a5315e90fce6d0abeac652.json b/backend/.sqlx/query-ebd8349ebda8ecb57e15257233a7ac50acac152fc0a5315e90fce6d0abeac652.json new file mode 100644 index 0000000000..55dc29c3a2 --- /dev/null +++ b/backend/.sqlx/query-ebd8349ebda8ecb57e15257233a7ac50acac152fc0a5315e90fce6d0abeac652.json @@ -0,0 +1,67 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT\n job_kind AS \"job_kind!: JobKind\",\n script_hash,\n flow_status AS \"flow_status!: Json>\",\n raw_flow AS \"raw_flow: Json>\"\n FROM v2_queue WHERE id = $1 AND workspace_id = $2 LIMIT 1", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_kind!: JobKind", + "type_info": { + "Custom": { + "name": "job_kind", + "kind": { + "Enum": [ + "script", + "preview", + "flow", + "dependencies", + "flowpreview", + "script_hub", + "identity", + "flowdependencies", + "http", + "graphql", + "postgresql", + "noop", + "appdependencies", + "deploymentcallback", + "singlescriptflow", + "flowscript", + "flownode", + "appscript" + ] + } + } + } + }, + { + "ordinal": 1, + "name": "script_hash", + "type_info": "Int8" + }, + { + "ordinal": 2, + "name": "flow_status!: Json>", + "type_info": "Jsonb" + }, + { + "ordinal": 3, + "name": "raw_flow: Json>", + "type_info": "Jsonb" + } + ], + "parameters": { + "Left": [ + "Uuid", + "Text" + ] + }, + "nullable": [ + true, + true, + true, + true + ] + }, + "hash": "ebd8349ebda8ecb57e15257233a7ac50acac152fc0a5315e90fce6d0abeac652" +} diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index 200344ddcc..b30df055bf 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -8,7 +8,6 @@ use axum::body::Body; use axum::http::HeaderValue; -use futures::TryFutureExt; use itertools::Itertools; use quick_cache::sync::Cache; use serde_json::value::RawValue; @@ -793,38 +792,26 @@ impl<'a> GetQuery<'a> { async fn resolve_raw_values( &self, db: &DB, - id: Uuid, kind: JobKind, - hash: Option, + runnable_id: Option, job: &mut JobExtended, ) { - let (raw_code, raw_lock, raw_flow) = ( - job.raw_code.take(), - job.raw_lock.take(), - job.raw_flow.take(), - ); - if self.with_flow { - // Try to fetch the flow from the cache, fallback to the preview flow. - // NOTE: This could check for the job kinds instead of the `or_else` but it's not - // necessary as `fetch_flow` return early if the job kind is not a preview one. - cache::job::fetch_flow(db, kind, hash) - .or_else(|_| cache::job::fetch_preview_flow(db, &id, raw_flow)) + job.raw_flow = match job.raw_flow.take() { + _ if !self.with_flow => None, + raw_flow => cache::job::fetch_flow(db, kind, runnable_id, raw_flow) .await .ok() - .inspect(|data| job.raw_flow = Some(sqlx::types::Json(data.raw_flow.clone()))); - } - if self.with_code { - // Try to fetch the code from the cache, fallback to the preview code. - // NOTE: This could check for the job kinds instead of the `or_else` but it's not - // necessary as `fetch_script` return early if the job kind is not a preview one. - cache::job::fetch_script(db, kind, hash) - .or_else(|_| cache::job::fetch_preview_script(db, &id, raw_lock, raw_code)) + .map(|data| sqlx::types::Json(data.raw_flow.clone())), + }; + (job.raw_code, job.raw_lock) = match (job.raw_code.take(), job.raw_lock.take()) { + _ if !self.with_code => (None, None), + (Some(raw_code), raw_lock) => (Some(raw_code), raw_lock), + (_, _) => cache::job::fetch_script(db, kind, runnable_id) .await .ok() - .inspect(|data| { - (job.raw_lock, job.raw_code) = (data.lock.clone(), Some(data.code.clone())) - }); - } + .map(|data| (Some(data.code.clone()), data.lock.clone())) + .unwrap_or_default(), + }; } async fn fetch_queued( @@ -847,7 +834,7 @@ impl<'a> GetQuery<'a> { self.check_auth(job.as_ref().map(|job| job.created_by.as_str()))?; if let Some(job) = job.as_mut() { - self.resolve_raw_values(db, job.id, job.job_kind, job.script_hash, job) + self.resolve_raw_values(db, job.job_kind, job.script_hash.map(|x| x.0), job) .await; } if self.with_flow { @@ -879,7 +866,7 @@ impl<'a> GetQuery<'a> { self.check_auth(cjob.as_ref().map(|job| job.created_by.as_str()))?; if let Some(job) = cjob.as_mut() { - self.resolve_raw_values(db, job.id, job.job_kind, job.script_hash, job) + self.resolve_raw_values(db, job.job_kind, job.script_hash.map(|x| x.0), job) .await; } if self.with_flow { diff --git a/backend/windmill-api/src/slack_approvals.rs b/backend/windmill-api/src/slack_approvals.rs index fcaa5b121c..3343fa2328 100644 --- a/backend/windmill-api/src/slack_approvals.rs +++ b/backend/windmill-api/src/slack_approvals.rs @@ -22,7 +22,6 @@ use windmill_common::{ cache, error::{self, Error}, jobs::JobKind, - scripts::ScriptHash, variables::{build_crypt, decrypt_value_with_mc}, }; @@ -978,7 +977,7 @@ async fn get_modal_blocks( let (job_kind, script_hash, raw_flow, parent_job_id, created_at, created_by, script_path, args) = sqlx::query!( "SELECT v2_queue.job_kind AS \"job_kind!: JobKind\", - v2_queue.script_hash AS \"script_hash: ScriptHash\", + v2_queue.script_hash, v2_queue.raw_flow AS \"raw_flow: sqlx::types::Json>\", v2_completed_job.parent_job AS \"parent_job: Uuid\", v2_completed_job.created_at AS \"created_at!: chrono::NaiveDateTime\", @@ -998,17 +997,11 @@ async fn get_modal_blocks( .ok_or_else(|| error::Error::BadRequest("This workflow is no longer running and has either already timed out or been cancelled or completed.".to_string())) .map(|r| (r.job_kind, r.script_hash, r.raw_flow, r.parent_job, r.created_at, r.created_by, r.script_path, r.args))?; - let flow_data = match cache::job::fetch_flow(&db, job_kind, script_hash).await { + let flow_data = match cache::job::fetch_flow(&db, job_kind, script_hash, raw_flow).await { Ok(data) => data, - Err(_) => { - if let Some(parent_job_id) = parent_job_id.as_ref() { - cache::job::fetch_preview_flow(&db, parent_job_id, raw_flow).await? - } else { - return Err(error::Error::BadRequest( - "This workflow is no longer running and has either already timed out or been cancelled or completed.".to_string(), - )); - } - } + Err(_) => return Err(error::Error::BadRequest( + "This workflow is no longer running and has either already timed out or been cancelled or completed.".to_string(), + )) }; let flow_value = &flow_data.flow; diff --git a/backend/windmill-common/src/cache.rs b/backend/windmill-common/src/cache.rs index 13d71dd402..a6c78394a4 100644 --- a/backend/windmill-common/src/cache.rs +++ b/backend/windmill-common/src/cache.rs @@ -302,6 +302,26 @@ pub enum RawData { Script(Arc), } +impl RawData { + pub fn from_raw( + raw_code: Option, + raw_lock: Option, + raw_flow: Option>>, + ) -> error::Result> { + match (raw_flow, raw_code, raw_lock) { + (Some(Json(raw_flow)), _, _) => FlowData::from_raw(raw_flow) + .map(Arc::new) + .map(Self::Flow) + .map(Some), + (_, Some(code), lock) => Ok(ScriptData { code, lock }) + .map(Arc::new) + .map(Self::Script) + .map(Some), + _ => Ok(None), + } + } +} + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct ScriptMetadata { pub language: Option, @@ -353,7 +373,6 @@ pub fn clear() { flow::clear(); script::clear(); app::clear(); - job::clear(); } pub mod flow { @@ -601,114 +620,19 @@ pub mod job { use super::*; use crate::jobs::JobKind; - #[cfg(not(feature = "scoped_cache"))] - lazy_static! { - /// Very small in-memory cache for "preview" jobs raw data. - static ref PREVIEWS: Cache = Cache::new(50); - } - - #[cfg(feature = "scoped_cache")] - lazy_static! { - /// Very small in-memory cache for "preview" jobs raw data. - static ref PREVIEWS: Cache<(ThreadId, Uuid), RawData> = Cache::new(50); - } - - /// Clear the job cache. - pub fn clear() { - PREVIEWS.clear(); - } - - #[track_caller] - pub fn fetch_preview_flow<'a, 'c>( - e: impl PgExecutor<'c> + 'a, - job: &'a Uuid, - // original raw values from `queue` or `completed_job` tables: - // kept for backward compatibility. - raw_flow: Option>>, - ) -> impl Future>> + 'a { - let fetch_preview = fetch_preview(e, job, None, None, raw_flow); - async move { - fetch_preview.await.and_then(|data| match data { - RawData::Flow(data) => Ok(data), - RawData::Script(_) => Err(error::Error::InternalErr(format!( - "Job ({job}) isn't a flow job." - ))), - }) - } - } - - #[track_caller] - pub fn fetch_preview_script<'a, 'c>( - e: impl PgExecutor<'c> + 'a, - job: &'a Uuid, - // original raw values from `queue` or `completed_job` tables: - // kept for backward compatibility. - raw_lock: Option, - raw_code: Option, - ) -> impl Future>> + 'a { - let fetch_preview = fetch_preview(e, job, raw_lock, raw_code, None); - async move { - fetch_preview.await.and_then(|data| match data { - RawData::Script(data) => Ok(data), - RawData::Flow(_) => Err(error::Error::InternalErr(format!( - "Job ({job}) isn't a script job." - ))), - }) - } - } - - #[track_caller] - pub fn fetch_preview<'a, 'c>( - e: impl PgExecutor<'c> + 'a, - job: &'a Uuid, - // original raw values from `queue` or `completed_job` tables: - // kept for backward compatibility. - raw_lock: Option, - raw_code: Option, - raw_flow: Option>>, - ) -> impl Future> + 'a { - let loc = Location::caller(); - let fetch = async move { - match (raw_lock, raw_code, raw_flow) { - (None, None, None) => sqlx::query!( - "SELECT raw_code, raw_lock, raw_flow AS \"raw_flow: Json>\" \ - FROM job WHERE id = $1 LIMIT 1", - job - ) - .fetch_optional(e) - .await - .map_err(Into::into) - .and_then(unwrap_or_error(&loc, "Preview", job)) - .map(|r| (r.raw_lock, r.raw_code, r.raw_flow)), - (lock, code, flow) => Ok((lock, code, flow)), - } - .and_then(|(lock, code, flow)| match flow { - Some(Json(flow)) => FlowData::from_raw(flow).map(Arc::new).map(RawData::Flow), - _ => Ok(RawData::Script(Arc::new(ScriptData { - code: code.unwrap_or_default(), - lock, - }))), - }) - }; - #[cfg(not(feature = "scoped_cache"))] - return PREVIEWS.get_or_insert_async(job, fetch); - #[cfg(feature = "scoped_cache")] - async move { - let job = &(std::thread::current().id(), job.clone()); - PREVIEWS.get_or_insert_async(job, fetch).await - } - } + // TODO(uael): new PREVIEWS cache #[track_caller] pub fn fetch_script<'c>( e: impl PgExecutor<'c>, kind: JobKind, - hash: Option, + runnable_id: Option, + // TODO(uael): add raw values here ) -> impl Future>> { use JobKind::*; let loc = Location::caller(); async move { - match (kind, hash.map(|ScriptHash(id)| id)) { + match (kind, runnable_id) { (FlowScript, Some(id)) => flow::fetch_script(e, FlowNodeId(id)).await, (Script | Dependencies, Some(hash)) => script::fetch(e, ScriptHash(hash)) .await @@ -727,16 +651,20 @@ pub mod job { pub fn fetch_flow<'c>( e: impl PgExecutor<'c> + Copy, kind: JobKind, - hash: Option, + runnable_id: Option, + raw_flow: Option>>, ) -> impl Future>> { use JobKind::*; let loc = Location::caller(); async move { - match (kind, hash.map(|ScriptHash(id)| id)) { + if let Some(Json(raw_flow)) = raw_flow { + return FlowData::from_raw(raw_flow).map(Arc::new); + } + match (kind, runnable_id) { (FlowDependencies, Some(id)) => flow::fetch_version(e, id).await, (FlowNode, Some(id)) => flow::fetch_flow(e, FlowNodeId(id)).await, (Flow, Some(id)) => match flow::fetch_version_lite(e, id).await { - Ok(raw_flow) => Ok(raw_flow), + Ok(data) => Ok(data), Err(_) => flow::fetch_version(e, id).await, }, _ => Err(error::Error::InternalErr(format!( diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 11e4f961e0..17692ef93b 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -29,12 +29,12 @@ use windmill_audit::ActionKind; use windmill_common::utils::now_from_db; use windmill_common::{ auth::{fetch_authed_from_permissioned_as, permissioned_as_to_username}, - cache::{self, FlowData}, + cache, db::{Authed, UserDB}, error::{self, to_anyhow, Error}, flow_status::{ - BranchAllStatus, FlowCleanupModule, FlowStatus, FlowStatusModule, FlowStatusModuleWParent, - Iterator, JobResult, RestartedFrom, RetryStatus, MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL, + BranchAllStatus, FlowStatus, FlowStatusModule, FlowStatusModuleWParent, Iterator, + JobResult, RestartedFrom, RetryStatus, MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL, }, flows::{ add_virtual_items_if_necessary, FlowModule, FlowModuleValue, FlowValue, InputTransform, @@ -868,7 +868,9 @@ pub async fn add_completed_job( #[cfg(feature = "enterprise")] if !success { async fn has_failure_module(db: &Pool, job: &QueuedJob) -> bool { - if let Ok(flow) = cache::job::fetch_flow(db, job.job_kind, job.script_hash).await { + if let Ok(flow) = + cache::job::fetch_flow(db, job.job_kind, job.script_hash.map(|x| x.0), None).await + { return flow.value().failure_module.is_some(); } sqlx::query_scalar!( @@ -3335,41 +3337,20 @@ pub async fn push<'c, 'd>( JobPayload::RawFlow { mut value, path, restarted_from } => { add_virtual_items_if_necessary(&mut value.modules); - let flow_status: FlowStatus = match restarted_from { - Some(restarted_from_val) => { - let (_, _, _, step_n, truncated_modules, user_states, cleanup_module) = - restarted_flows_resolution( - _db, - workspace_id, - restarted_from_val.flow_job_id, - restarted_from_val.step_id.as_str(), - restarted_from_val.branch_or_iteration_n, - ) - .await?; - FlowStatus { - step: step_n, - modules: truncated_modules, - // failure_module is reset - failure_module: Box::new(FlowStatusModuleWParent { - parent_module: None, - module_status: FlowStatusModule::WaitingForPriorSteps { - id: "failure".to_string(), - }, - }), - cleanup_module, - // retry status is reset - retry: RetryStatus { fail_count: 0, failed_jobs: vec![] }, - // TODO: for now, flows with approval conditions aren't supported for restart - approval_conditions: None, - restarted_from: Some(RestartedFrom { - flow_job_id: restarted_from_val.flow_job_id, - step_id: restarted_from_val.step_id, - branch_or_iteration_n: restarted_from_val.branch_or_iteration_n, - }), - user_states, - preprocessor_module: None, - } - } + let flow_status = match restarted_from { + Some(restarted_from) => restarted_flows_resolution( + workspace_id, + &value, + sqlx::query_scalar!( + "SELECT flow_status AS \"flow_status: Json>\" + FROM v2_job_completed WHERE id = $1", + restarted_from.flow_job_id, + ) + .fetch_optional(_db) + .await? + .flatten(), + restarted_from, + )?, _ => { value.preprocessor_module = None; FlowStatus::new(&value) @@ -3548,45 +3529,26 @@ pub async fn push<'c, 'd>( ) } JobPayload::RestartedFlow { completed_job_id, step_id, branch_or_iteration_n } => { - let ( - version, - flow_path, - flow_data, - step_n, - truncated_modules, - user_states, - cleanup_module, - ) = restarted_flows_resolution( - _db, - workspace_id, - completed_job_id, - step_id.as_str(), - branch_or_iteration_n, + let (job_kind, runnable_id, runnable_path, flow_status) = sqlx::query!( + "SELECT kind AS \"kind: JobKind\", runnable_id, runnable_path, + flow_status AS \"flow_status: Json>\" + FROM v2_job JOIN v2_job_completed USING (id) + WHERE id = $1", + completed_job_id ) - .await?; - let restarted_flow_status = FlowStatus { - step: step_n, - modules: truncated_modules, - // failure_module is reset - failure_module: Box::new(FlowStatusModuleWParent { - parent_module: None, - module_status: FlowStatusModule::WaitingForPriorSteps { - id: "failure".to_string(), - }, - }), - cleanup_module, - // retry status is reset - retry: RetryStatus { fail_count: 0, failed_jobs: vec![] }, - // TODO: for now, flows with approval conditions aren't supported for restart - approval_conditions: None, - restarted_from: Some(RestartedFrom { - flow_job_id: completed_job_id, - step_id, - branch_or_iteration_n, - }), - user_states, - preprocessor_module: None, - }; + .map(|r| (r.kind, r.runnable_id, r.runnable_path, r.flow_status)) + .fetch_optional(_db) + .await? + .ok_or_else(|| { + Error::InternalErr(format!("{:?}: completed job not found", completed_job_id)) + })?; + let flow_data = cache::job::fetch_flow(_db, job_kind, runnable_id, None).await?; + let flow_status = restarted_flows_resolution( + workspace_id, + flow_data.value(), + flow_status, + RestartedFrom { flow_job_id: completed_job_id, step_id, branch_or_iteration_n }, + )?; let value = flow_data.value(); let priority = value.priority; let concurrency_key = value.concurrency_key.clone(); @@ -3595,19 +3557,19 @@ pub async fn push<'c, 'd>( let cache_ttl = value.cache_ttl.map(|x| x as i32); // Keep inserting `value` if not all workers are updated. // Starting at `v1.440`, the value is fetched on pull from the version id. - let value_o = if version.is_none() || !*MIN_VERSION_IS_AT_LEAST_1_440.read().await { + let value_o = if runnable_id.is_none() || !*MIN_VERSION_IS_AT_LEAST_1_440.read().await { Some(value.clone()) } else { // `raw_flow` is fetched on pull. None }; ( - version, - flow_path, + runnable_id, + runnable_path, None, JobKind::Flow, value_o, - Some(restarted_flow_status), + Some(flow_status), None, concurrency_key, concurrent_limit, @@ -3986,50 +3948,19 @@ pub fn canceled_job_to_result(job: &QueuedJob) -> serde_json::Value { serde_json::json!({"message": format!("Job canceled: {reason} by {canceler}"), "name": "Canceled", "reason": reason, "canceler": canceler}) } -async fn restarted_flows_resolution( - db: &Pool, +fn restarted_flows_resolution( workspace_id: &str, - completed_flow_id: Uuid, - restart_step_id: &str, - branch_or_iteration_n: Option, -) -> Result< - ( - Option, - Option, - Arc, - i32, - Vec, - HashMap, - FlowCleanupModule, - ), - Error, -> { - let row = sqlx::query!( - "SELECT - script_path, script_hash AS \"script_hash: ScriptHash\", - job_kind AS \"job_kind!: JobKind\", - flow_status AS \"flow_status: Json>\", - raw_flow AS \"raw_flow: Json>\" - FROM v2_completed_job WHERE id = $1 and workspace_id = $2", - completed_flow_id, - workspace_id, - ) - .fetch_one(db) // TODO: should we try to use the passed-in `tx` here? - .await - .map_err(|err| { - Error::InternalErr(format!( - "completed job not found for UUID {} in workspace {}: {}", - completed_flow_id, workspace_id, err - )) - })?; + flow_value: &FlowValue, + flow_status: Option>>, + restart_from: RestartedFrom, +) -> error::Result { + let RestartedFrom { + flow_job_id: completed_flow_id, + step_id: restart_step_id, + branch_or_iteration_n, + } = &restart_from; - let flow_data = cache::job::fetch_flow(db, row.job_kind, row.script_hash) - .or_else(|_| cache::job::fetch_preview_flow(db, &completed_flow_id, row.raw_flow)) - .await?; - let flow_value = flow_data.value(); - let flow_status = row - .flow_status - .as_ref() + let mut flow_status = flow_status .and_then(|v| serde_json::from_str::(v.get()).ok()) .ok_or(Error::InternalErr(format!( "Unable to parse flow status for job {} in workspace {}", @@ -4040,20 +3971,21 @@ async fn restarted_flows_resolution( let mut dependent_module = false; let mut truncated_modules: Vec = vec![]; for module in flow_status.modules { + let id = module.id(); let Some(module_definition) = flow_value .modules .iter() - .find(|flow_value_module| flow_value_module.id == module.id()) + .find(|flow_value_module| &flow_value_module.id == &id) else { // skip module as it doesn't appear in the flow_value anymore continue; }; - if module.id() == restart_step_id { + if &id == restart_step_id { // if the module ID is the one we want to restart the flow at, or if it's past it in the flow, // set the module as WaitingForPriorSteps as it needs to be re-run if branch_or_iteration_n.is_none() || branch_or_iteration_n.unwrap() == 0 { // The module as WaitingForPriorSteps as the entire module (i.e. all the branches) need to be re-run - truncated_modules.push(FlowStatusModule::WaitingForPriorSteps { id: module.id() }); + truncated_modules.push(FlowStatusModule::WaitingForPriorSteps { id }); } else { // expect a module to be either a branchall (resp. loop), and resume the flow from this branch (resp. iteration) let branch_or_iteration_n = branch_or_iteration_n.unwrap(); @@ -4082,7 +4014,7 @@ async fn restarted_flows_resolution( new_flow_jobs_success.truncate(branch_or_iteration_n); } truncated_modules.push(FlowStatusModule::InProgress { - id: module.id(), + id, job: new_flow_jobs[new_flow_jobs.len() - 1], // set to last finished job from completed flow iterator: None, flow_jobs: Some(new_flow_jobs), @@ -4120,7 +4052,7 @@ async fn restarted_flows_resolution( new_flow_jobs_success.truncate(branch_or_iteration_n); } truncated_modules.push(FlowStatusModule::InProgress { - id: module.id(), + id, job: new_flow_jobs[new_flow_jobs.len() - 1], // set to last finished job from completed flow iterator: Some(Iterator { index: branch_or_iteration_n - 1, // same deal as above, this refers to the last finished job @@ -4146,7 +4078,7 @@ async fn restarted_flows_resolution( } dependent_module = true; } else if dependent_module { - truncated_modules.push(FlowStatusModule::WaitingForPriorSteps { id: module.id() }); + truncated_modules.push(FlowStatusModule::WaitingForPriorSteps { id }); } else { // else we simply "transfer" the module from the completed flow to the new one if it's a success step_n = step_n + 1; @@ -4167,13 +4099,19 @@ async fn restarted_flows_resolution( ))); } - Ok(( - row.script_hash.map(|x| x.0), - row.script_path, - flow_data, - step_n, - truncated_modules, - flow_status.user_states, - flow_status.cleanup_module, - )) + flow_status.modules = truncated_modules; + flow_status.step = step_n; + flow_status.restarted_from = Some(restart_from); + // failure_module is reset + flow_status.failure_module = Box::new(FlowStatusModuleWParent { + parent_module: None, + module_status: FlowStatusModule::WaitingForPriorSteps { id: "failure".to_string() }, + }); + // retry status is reset + flow_status.retry = RetryStatus { fail_count: 0, failed_jobs: vec![] }; + // no preprocessor module + flow_status.preprocessor_module = None; + // TODO: for now, flows with approval conditions aren't supported for restart + flow_status.approval_conditions = None; + Ok(flow_status) } diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 3ea0c8109d..94015812ce 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -1974,22 +1974,9 @@ async fn handle_queued_job( } let started = Instant::now(); - // Pre-fetch preview jobs raw values if necessary. - // The `raw_*` values passed to this function are the original raw values from `queue` tables, - // they are kept for backward compatibility as they have been moved to the `job` table. - let preview_data = match (job.job_kind, job.script_hash) { - ( - JobKind::Preview - | JobKind::Dependencies - | JobKind::FlowPreview - | JobKind::Flow - | JobKind::FlowDependencies, - None, - ) => Some(cache::job::fetch_preview(db, &job.id, raw_lock, raw_code, raw_flow).await?), - _ => None, - }; + let raw_data = RawData::from_raw(raw_code, raw_lock, raw_flow)?; let cached_res_path = if job.cache_ttl.is_some() { - Some(cached_result_path(db, &client.get_authed().await, &job, preview_data.as_ref()).await) + Some(cached_result_path(db, &client.get_authed().await, &job, raw_data.as_ref()).await) } else { None }; @@ -2030,10 +2017,11 @@ async fn handle_queued_job( } }; if job.is_flow() { - let flow_data = match preview_data { + let runnable_id = job.script_hash.map(|x| x.0); + let flow_data = match raw_data { Some(RawData::Flow(data)) => data, // Not a preview: fetch from the cache or the database. - _ => cache::job::fetch_flow(db, job.job_kind, job.script_hash).await?, + _ => cache::job::fetch_flow(db, job.job_kind, runnable_id, None).await?, }; handle_flow( job, @@ -2091,7 +2079,7 @@ async fn handle_queued_job( JobKind::Dependencies => { handle_dependency_job( &job, - preview_data.as_ref(), + raw_data.as_ref(), &mut mem_peak, &mut canceled_by, job_dir, @@ -2107,7 +2095,7 @@ async fn handle_queued_job( JobKind::FlowDependencies => { handle_flow_dependency_job( &job, - preview_data.as_ref(), + raw_data.as_ref(), &mut mem_peak, &mut canceled_by, job_dir, @@ -2143,13 +2131,13 @@ async fn handle_queued_job( .unwrap_or_else(|| serde_json::from_str("{}").unwrap())), _ => { let metric_timer = Instant::now(); - let preview_data = preview_data.and_then(|data| match data { + let raw_data = raw_data.and_then(|data| match data { RawData::Script(data) => Some(data), _ => None, }); let r = handle_code_execution_job( job.as_ref(), - preview_data, + raw_data, db, client, job_dir, diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 2805b27c88..adbd263aa9 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -11,6 +11,9 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; +#[cfg(feature = "benchmark")] +use crate::bench::BenchmarkIter; + use crate::common::{cached_result_path, save_in_cache}; use crate::js_eval::{eval_timeout, IdContext}; use crate::{ @@ -18,7 +21,6 @@ use crate::{ KEEP_JOB_DIR, }; use anyhow::Context; -use futures::TryFutureExt; use mappable_rc::Marc; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; @@ -30,8 +32,6 @@ use tracing::instrument; use uuid::Uuid; use windmill_common::add_time; use windmill_common::auth::JobPerms; -#[cfg(feature = "benchmark")] -use crate::bench::BenchmarkIter; use windmill_common::cache::{self, RawData}; use windmill_common::db::Authed; use windmill_common::flow_status::{ @@ -42,7 +42,6 @@ use windmill_common::jobs::{ script_hash_to_tag_and_limits, script_path_to_payload, BranchResults, JobKind, JobPayload, OnBehalfOf, QueuedJob, RawCode, ENTRYPOINT_OVERRIDE, }; -use windmill_common::scripts::ScriptHash; use windmill_common::users::username_to_permissioned_as; use windmill_common::utils::WarnAfterExt; use windmill_common::worker::to_raw_value; @@ -221,7 +220,7 @@ pub async fn update_flow_status_after_job_completion_internal( let (job_kind, script_hash, old_status, raw_flow) = sqlx::query!( "SELECT job_kind AS \"job_kind!: JobKind\", - script_hash AS \"script_hash: ScriptHash\", + script_hash, flow_status AS \"flow_status!: Json>\", raw_flow AS \"raw_flow: Json>\" FROM v2_queue WHERE id = $1 AND workspace_id = $2 LIMIT 1", @@ -248,9 +247,7 @@ pub async fn update_flow_status_after_job_completion_internal( )) })?; - let flow_data = cache::job::fetch_flow(db, job_kind, script_hash) - .or_else(|_| cache::job::fetch_preview_flow(db, &flow, raw_flow)) - .await?; + let flow_data = cache::job::fetch_flow(db, job_kind, script_hash, raw_flow).await?; let flow_value = flow_data.value(); let module_step = Step::from_i32_and_len(old_status.step, old_status.modules.len());