From d7b688d3529656a0918faba3c93d65cb6977ee1f Mon Sep 17 00:00:00 2001
From: Lucas Abel <22837557+uael@users.noreply.github.com>
Date: Thu, 14 Nov 2024 08:49:12 +0100
Subject: [PATCH] nit: generalize usage of `has_failure_module` (#4706)

---
 backend/windmill-common/src/flows.rs       | 27 ++++++++++++++++++++++
 backend/windmill-queue/src/jobs.rs         | 20 ++++------------
 backend/windmill-worker/src/worker_flow.rs | 21 ++---------------
 3 files changed, 33 insertions(+), 35 deletions(-)

diff --git a/backend/windmill-common/src/flows.rs b/backend/windmill-common/src/flows.rs
index d357a5ec72..5d7e1b4621 100644
--- a/backend/windmill-common/src/flows.rs
+++ b/backend/windmill-common/src/flows.rs
@@ -16,6 +16,7 @@ use rand::Rng;
 use serde::{Deserialize, Serialize, Serializer};
 
 use crate::{
+    error::Error,
     more_serde::{default_empty_string, default_id, default_null, default_true, is_default},
     scripts::{Schema, ScriptHash, ScriptLang},
 };
@@ -651,3 +652,29 @@ pub fn add_virtual_items_if_necessary(modules: &mut Vec<FlowModule>) {
         });
     }
 }
+
+pub async fn has_failure_module<'c>(flow: sqlx::types::Uuid, db: &sqlx::Pool<sqlx::Postgres>, completed: bool) -> Result<bool, Error> {
+    if completed {
+        sqlx::query_scalar!(
+            "SELECT raw_flow->'failure_module' != 'null'::jsonb
+            FROM completed_job
+            WHERE id = $1",
+            flow
+        )
+    } else {
+        sqlx::query_scalar!(
+            "SELECT raw_flow->'failure_module' != 'null'::jsonb
+            FROM queue
+            WHERE id = $1",
+            flow
+        )
+    }
+    .fetch_one(db)
+    .await
+    .map_err(|e| {
+        Error::InternalErr(format!(
+            "error during retrieval of has_failure_module: {e:#}"
+        ))
+    })
+    .map(|v| v.unwrap_or(false))
+}
diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs
index 9b94c0f198..275034efe7 100644
--- a/backend/windmill-queue/src/jobs.rs
+++ b/backend/windmill-queue/src/jobs.rs
@@ -71,6 +71,9 @@ use windmill_common::BASE_URL;
 #[cfg(feature = "cloud")]
 use windmill_common::users::SUPERADMIN_SYNC_EMAIL;
 
+#[cfg(feature = "enterprise")]
+use windmill_common::flows::has_failure_module;
+
 #[cfg(feature = "enterprise")]
 use windmill_common::worker::CLOUD_HOSTED;
 
@@ -482,13 +485,6 @@ where
     }
 }
 
-#[cfg(feature = "enterprise")]
-#[derive(Deserialize)]
-struct RawFlowFailureModule {
-    #[cfg(feature = "enterprise")]
-    failure_module: Option<Box<RawValue>>,
-}
-
 #[instrument(level = "trace", skip_all)]
 pub async fn add_completed_job_error<R: rsmq_async::RsmqConnection + Clone + Send>(
     db: &Pool<Postgres>,
@@ -952,15 +948,7 @@ pub async fn add_completed_job<
         } else if !skip_downstream_error_handlers
             && (matches!(queued_job.job_kind, JobKind::Script)
                 || matches!(queued_job.job_kind, JobKind::Flow)
-                    && queued_job
-                        .raw_flow
-                        .as_ref()
-                        .and_then(|v| {
-                            serde_json::from_str::<RawFlowFailureModule>((**v).get())
-                                .ok()
-                                .and_then(|v| v.failure_module)
-                        })
-                        .is_none())
+                    && !has_failure_module(job_id, db, true).await.unwrap_or(false))
             && queued_job.parent_job.is_none()
         {
             let result = serde_json::from_str(
diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs
index a1a4fd9904..5421b502e4 100644
--- a/backend/windmill-worker/src/worker_flow.rs
+++ b/backend/windmill-worker/src/worker_flow.rs
@@ -49,7 +49,7 @@ use windmill_common::{
         Approval, BranchAllStatus, BranchChosen, FlowStatus, FlowStatusModule, RetryStatus,
         MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL,
     },
-    flows::{FlowModule, FlowModuleValue, FlowValue, InputTransform, Retry, Suspend},
+    flows::{has_failure_module, FlowModule, FlowModuleValue, FlowValue, InputTransform, Retry, Suspend},
 };
 use windmill_queue::schedule::get_schedule_opt;
 use windmill_queue::{
@@ -963,7 +963,7 @@ pub async fn update_flow_status_after_job_completion_internal<
             false
                 if !is_failure_step
                     && !skip_error_handler
-                    && has_failure_module(flow, db).await? =>
+                    && has_failure_module(flow, db, false).await? =>
             {
                 true
             }
@@ -1290,23 +1290,6 @@ async fn compute_skip_branchall_failure<'c>(
         }))
 }
 
-async fn has_failure_module<'c>(flow: Uuid, db: &DB) -> Result<bool, Error> {
-    sqlx::query_scalar::<_, Option<bool>>(
-        "SELECT raw_flow->'failure_module' != 'null'::jsonb
-        FROM queue
-        WHERE id = $1",
-    )
-    .bind(flow)
-    .fetch_one(db)
-    .await
-    .map_err(|e| {
-        Error::InternalErr(format!(
-            "error during retrieval of has_failure_module: {e:#}"
-        ))
-    })
-    .map(|v| v.unwrap_or(false))
-}
-
 // async fn retrieve_cleanup_module<'c>(flow_uuid: Uuid, db: &DB) -> Result<FlowCleanupModule, Error> {
 //     tracing::warn!("Retrieving cleanup module of flow {}", flow_uuid);
 //     let raw_value = sqlx::query_scalar!(