From c7a327e1d8a1b873cf3d7ee2edddcef05484a004 Mon Sep 17 00:00:00 2001 From: yacine Bouraroui Date: Tue, 1 Oct 2024 12:30:11 +0200 Subject: [PATCH] resuable raw flow parsing logic. 2- enforced some queries to be compile-checked. some types "masking" and a better alternative to some 'select * from queue' --- backend/src/monitor.rs | 30 +++++++++------ backend/windmill-api/src/resources.rs | 32 +++++++++++++--- backend/windmill-common/src/flow_status.rs | 14 +++++++ backend/windmill-common/src/flows.rs | 14 +++++++ backend/windmill-common/src/jobs.rs | 43 ++++++++++------------ backend/windmill-queue/src/jobs.rs | 25 +++++++++++-- backend/windmill-worker/src/worker_flow.rs | 4 +- 7 files changed, 115 insertions(+), 47 deletions(-) diff --git a/backend/src/monitor.rs b/backend/src/monitor.rs index 92f69a0d64..ea7947c31f 100644 --- a/backend/src/monitor.rs +++ b/backend/src/monitor.rs @@ -19,6 +19,7 @@ use tokio::{ sync::{mpsc, RwLock}, }; +use uuid::Uuid; #[cfg(feature = "embedding")] use windmill_api::embeddings::update_embeddings_db; use windmill_api::{ @@ -32,7 +33,7 @@ use windmill_common::{ auth::JWT_SECRET, ee::CriticalErrorChannel, error, - flow_status::FlowStatusModule, + flow_status::{FlowStatusModule, ParsedFlowStatusGetter as _}, global_settings::{ BASE_URL_SETTING, BUNFIG_INSTALL_SCOPES_SETTING, CRITICAL_ERROR_CHANNELS_SETTING, DEFAULT_TAGS_PER_WORKSPACE_SETTING, DEFAULT_TAGS_WORKSPACES_SETTING, @@ -1425,7 +1426,7 @@ async fn handle_zombie_flows( } ); report_critical_error(reason.clone(), db.clone()).await; - cancel_zombie_flow_job(db, flow, &rsmq, reason).await?; + cancel_zombie_flow_job(db, &flow.id, &flow.workspace_id, &rsmq, reason).await?; } } @@ -1441,11 +1442,17 @@ async fn handle_zombie_flows( .fetch_all(db) .await?; + #[derive(sqlx::FromRow, Debug)] + struct InQueueJobResult { + id: uuid::Uuid, + workspace_id: String, + } + for flow in flows2 { - let in_queue = sqlx::query_as::<_, QueuedJob>( - "SELECT * FROM queue WHERE id = $1 AND running = true AND canceled = false", + let in_queue = sqlx::query_as!(InQueueJobResult, + "SELECT id, workspace_id FROM queue WHERE id = $1 AND running = true AND canceled = false", + flow.parent_flow_id ) - .bind(flow.parent_flow_id) .fetch_optional(db) .await?; if let Some(job) = in_queue { @@ -1455,7 +1462,7 @@ async fn handle_zombie_flows( job.workspace_id, flow.last_ping ); - cancel_zombie_flow_job(db, job, &rsmq, + cancel_zombie_flow_job(db, &job.id, &job.workspace_id, &rsmq, format!("Flow {} cancelled as one of the parallel branch {} was unable to make the last transition ", flow.parent_flow_id, flow.job_id)) .await?; } else { @@ -1467,21 +1474,22 @@ async fn handle_zombie_flows( async fn cancel_zombie_flow_job( db: &Pool, - flow: QueuedJob, + job_id: &Uuid, + workspace_id: &str, rsmq: &Option, message: String, ) -> Result<(), error::Error> { let tx = db.begin().await.unwrap(); tracing::error!( "zombie flow detected: {} in workspace {}. Cancelling it.", - flow.id, - flow.workspace_id + job_id, + workspace_id ); let (ntx, _) = cancel_job( "monitor", Some(message), - flow.id, - flow.workspace_id.as_str(), + *job_id, + workspace_id, tx, db, rsmq.clone(), diff --git a/backend/windmill-api/src/resources.rs b/backend/windmill-api/src/resources.rs index 1fb5b12259..d1de3c4328 100644 --- a/backend/windmill-api/src/resources.rs +++ b/backend/windmill-api/src/resources.rs @@ -31,7 +31,6 @@ use windmill_audit::ActionKind; use windmill_common::{ db::UserDB, error::{Error, JsonResult, Result}, - jobs::QueuedJob, utils::{not_found_if_none, paginate, require_admin, Pagination, StripPath}, variables, }; @@ -58,7 +57,10 @@ pub fn workspaced_service() -> Router { .route("/type/exists/:name", get(exists_resource_type)) .route("/type/update/:name", post(update_resource_type)) .route("/type/delete/:name", delete(delete_resource_type)) - .route("/file_resource_type_to_file_ext_map", get(file_resource_ext_to_resource_type)) + .route( + "/file_resource_type_to_file_ext_map", + get(file_resource_ext_to_resource_type), + ) .route("/type/create", post(create_resource_type)) } @@ -533,11 +535,29 @@ pub async fn transform_json_value<'c>( } Value::String(y) if y.starts_with("$") && job_id.is_some() => { let mut tx = authed_transaction_or_default(authed, user_db.clone(), db).await?; - let job = sqlx::query_as::<_, QueuedJob>( - "SELECT * FROM queue WHERE id = $1 AND workspace_id = $2", + + #[derive(sqlx::FromRow, Debug)] + struct QueuedJobLite { + pub id: Uuid, + pub workspace_id: String, + pub parent_job: Option, + pub created_by: String, + pub email: String, + pub permissioned_as: String, + pub script_path: Option, + pub schedule_path: Option, + pub root_job: Option, + pub flow_step_id: Option, + pub scheduled_for: chrono::DateTime, + } + + let job = sqlx::query_as!( + QueuedJobLite, + "SELECT id, workspace_id,parent_job, created_by, email, permissioned_as, script_path, schedule_path, root_job, flow_step_id, scheduled_for + FROM queue WHERE id = $1 AND workspace_id = $2", + job_id.unwrap(), + workspace ) - .bind(job_id.unwrap()) - .bind(workspace) .fetch_optional(&mut *tx) .await?; tx.commit().await?; diff --git a/backend/windmill-common/src/flow_status.rs b/backend/windmill-common/src/flow_status.rs index 035cadecd8..fe5c1467e1 100644 --- a/backend/windmill-common/src/flow_status.rs +++ b/backend/windmill-common/src/flow_status.rs @@ -45,6 +45,20 @@ pub struct FlowStatus { pub restarted_from: Option, } +pub trait FlowStatusGetter { + fn get_raw_flow_status(&self) -> Option<&sqlx::types::Json>>; +} +pub trait ParsedFlowStatusGetter { + fn parse_flow_status(&self) -> Option; +} + +impl ParsedFlowStatusGetter for I { + fn parse_flow_status(&self) -> Option { + self.get_raw_flow_status() + .and_then(|v| serde_json::from_str::((**v).get()).ok()) + } +} + #[derive(Serialize, Deserialize, Debug, Clone, Default)] #[serde(default)] pub struct RetryStatus { diff --git a/backend/windmill-common/src/flows.rs b/backend/windmill-common/src/flows.rs index a5040a9da3..05efc6e136 100644 --- a/backend/windmill-common/src/flows.rs +++ b/backend/windmill-common/src/flows.rs @@ -124,6 +124,20 @@ pub struct FlowValue { pub concurrency_key: Option, } +pub trait FlowValueGetter { + fn get_raw_flow_value(&self) -> Option<&sqlx::types::Json>>; +} +pub trait ParsedFlowValueGetter { + fn parse_raw_flow(&self) -> Option; +} + +impl ParsedFlowValueGetter for I { + fn parse_raw_flow(&self) -> Option { + self.get_raw_flow_value() + .and_then(|v| serde_json::from_str::((**v).get()).ok()) + } +} + #[derive(Deserialize, Serialize, Debug, Clone)] pub struct StopAfterIf { pub expr: String, diff --git a/backend/windmill-common/src/jobs.rs b/backend/windmill-common/src/jobs.rs index f7cfe1c9ee..71a0fe8423 100644 --- a/backend/windmill-common/src/jobs.rs +++ b/backend/windmill-common/src/jobs.rs @@ -15,8 +15,8 @@ pub const PREPROCESSOR_FAKE_ENTRYPOINT: &str = "__WM_PREPROCESSOR"; use crate::{ error::{self, to_anyhow, Error}, - flow_status::{FlowStatus, RestartedFrom}, - flows::{FlowValue, Retry}, + flow_status::{FlowStatusGetter, RestartedFrom}, + flows::{FlowValue, FlowValueGetter, Retry}, get_latest_deployed_hash_for_path, scripts::{ScriptHash, ScriptLang}, worker::{to_raw_value, TMP_DIR}, @@ -119,10 +119,7 @@ pub struct QueuedJob { impl QueuedJob { pub fn script_path(&self) -> &str { - self.script_path - .as_ref() - .map(String::as_str) - .unwrap_or("tmp/main") + self.script_path.as_deref().unwrap_or("tmp/main") } pub fn is_flow(&self) -> bool { matches!( @@ -139,19 +136,17 @@ impl QueuedJob { self.script_path() ) } +} - pub fn parse_raw_flow(&self) -> Option { - self.raw_flow.as_ref().and_then(|v| { - let str = (**v).get(); - // tracing::error!("raw_flow: {}", str); - return serde_json::from_str::(str).ok(); - }) +impl FlowValueGetter for QueuedJob { + fn get_raw_flow_value(&self) -> Option<&sqlx::types::Json>> { + self.raw_flow.as_ref() } +} - pub fn parse_flow_status(&self) -> Option { - self.flow_status - .as_ref() - .and_then(|v| serde_json::from_str::((**v).get()).ok()) +impl FlowStatusGetter for QueuedJob { + fn get_raw_flow_status(&self) -> Option<&sqlx::types::Json>> { + self.flow_status.as_ref() } } @@ -269,17 +264,17 @@ impl CompletedJob { .map(|r| serde_json::from_str(r.get()).ok()) .flatten() } +} - pub fn parse_raw_flow(&self) -> Option { - self.raw_flow - .as_ref() - .and_then(|v| serde_json::from_str::((**v).get()).ok()) +impl FlowValueGetter for CompletedJob { + fn get_raw_flow_value(&self) -> Option<&sqlx::types::Json>> { + self.raw_flow.as_ref() } +} - pub fn parse_flow_status(&self) -> Option { - self.flow_status - .as_ref() - .and_then(|v| serde_json::from_str::((**v).get()).ok()) +impl FlowStatusGetter for CompletedJob { + fn get_raw_flow_status(&self) -> Option<&sqlx::types::Json>> { + self.flow_status.as_ref() } } diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 74383b6b82..0467272ef4 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -44,11 +44,13 @@ use windmill_common::{ 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, FlowCleanupModule, FlowStatus, FlowStatusGetter, FlowStatusModule, + FlowStatusModuleWParent, Iterator, JobResult, ParsedFlowStatusGetter, RestartedFrom, + RetryStatus, MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL, }, flows::{ add_virtual_items_if_necessary, FlowModule, FlowModuleValue, FlowValue, InputTransform, + ParsedFlowValueGetter, }, jobs::{ get_payload_tag_from_prefixed_path, CompletedJob, JobKind, JobPayload, QueuedJob, RawCode, @@ -2271,6 +2273,20 @@ pub async fn get_result_by_id( node_id: String, json_path: Option, ) -> error::Result> { + #[derive(sqlx::FromRow, Debug, Serialize, Clone)] + struct RunningFlowJobResult { + pub id: Uuid, + pub flow_status: Option>>, + } + + impl FlowStatusGetter for RunningFlowJobResult { + fn get_raw_flow_status( + &self, + ) -> Option<&sqlx::types::Json>> { + self.flow_status.as_ref() + } + } + match get_result_by_id_from_running_flow( &db, w_id.as_str(), @@ -2282,11 +2298,12 @@ pub async fn get_result_by_id( { Ok(res) => Ok(res), Err(_) => { - let running_flow_job =sqlx::query_as::<_, QueuedJob>( - "SELECT * FROM queue WHERE COALESCE((SELECT root_job FROM queue WHERE id = $1), $1) = id AND workspace_id = $2" + let running_flow_job =sqlx::query_as::<_, RunningFlowJobResult>( + "SELECT id, flow_status FROM queue WHERE COALESCE((SELECT root_job FROM queue WHERE id = $1), $1) = id AND workspace_id = $2" ).bind(flow_id) .bind(&w_id) .fetch_optional(&db).await?; + match running_flow_job { Some(job) => { let restarted_from = windmill_common::utils::not_found_if_none( diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index cf979fc340..a257e50ea9 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -35,9 +35,9 @@ use windmill_common::auth::JobPerms; use windmill_common::bench::BenchmarkIter; use windmill_common::db::Authed; use windmill_common::flow_status::{ - ApprovalConditions, FlowStatusModuleWParent, Iterator, JobResult, + ApprovalConditions, FlowStatusModuleWParent, Iterator, JobResult, ParsedFlowStatusGetter, }; -use windmill_common::flows::add_virtual_items_if_necessary; +use windmill_common::flows::{add_virtual_items_if_necessary, ParsedFlowValueGetter}; use windmill_common::jobs::{ script_hash_to_tag_and_limits, script_path_to_payload, BranchResults, JobPayload, QueuedJob, RawCode, ENTRYPOINT_OVERRIDE,