backend: re-work job table usages after v2

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<Box<RawValue>>\"\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<Box<RawValue>>",
+        "type_info": "Jsonb"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Uuid"
+      ]
+    },
+    "nullable": [
+      false,
+      true,
+      true,
+      true
+    ]
+  },
+  "hash": "111ee00fa6661fd4d3bb1d2e567fdaa1cbd3fe93b590b97c549ff2ad0016da1a"
+}

backend/.sqlx/query-4a4971ab285cc5ca0dc4026d75a31981cb91ccabdaa786ebadb6d80e8a8b38db.json

@@ -0,0 +1,22 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT flow_status AS \"flow_status: Json<Box<RawValue>>\"\n                            FROM v2_job_completed WHERE id = $1",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "flow_status: Json<Box<RawValue>>",
+        "type_info": "Jsonb"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Uuid"
+      ]
+    },
+    "nullable": [
+      true
+    ]
+  },
+  "hash": "4a4971ab285cc5ca0dc4026d75a31981cb91ccabdaa786ebadb6d80e8a8b38db"
+}

backend/.sqlx/query-99a8b3598bf9f63b5aed5fdbbe16d91cee8e3dcfad07b5fef9398e98f42d7ab2.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<Box<RawValue>>\",\n                raw_flow AS \"raw_flow: Json<Box<RawValue>>\"\n            FROM v2_as_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<Box<RawValue>>",
+        "type_info": "Jsonb"
+      },
+      {
+        "ordinal": 3,
+        "name": "raw_flow: Json<Box<RawValue>>",
+        "type_info": "Jsonb"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Uuid",
+        "Text"
+      ]
+    },
+    "nullable": [
+      true,
+      true,
+      true,
+      true
+    ]
+  },
+  "hash": "99a8b3598bf9f63b5aed5fdbbe16d91cee8e3dcfad07b5fef9398e98f42d7ab2"
+}

backend/.sqlx/query-f3d8c6804a4d3ebf69704ebe7956c38ca756fdc998cb1fe3e79f54269a914b49.json

@@ -0,0 +1,91 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT\n            v2_as_queue.job_kind AS \"job_kind!: JobKind\",\n            v2_as_queue.script_hash,\n            v2_as_queue.raw_flow AS \"raw_flow: sqlx::types::Json<Box<RawValue>>\",\n            v2_as_completed_job.parent_job AS \"parent_job: Uuid\",\n            v2_as_completed_job.created_at AS \"created_at!: chrono::NaiveDateTime\",\n            v2_as_completed_job.created_by AS \"created_by!\",\n            v2_as_queue.script_path,\n            v2_as_queue.args AS \"args: sqlx::types::Json<Box<RawValue>>\"\n        FROM v2_as_queue\n        JOIN v2_as_completed_job ON v2_as_completed_job.parent_job = v2_as_queue.id\n        WHERE v2_as_completed_job.id = $1 AND v2_as_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<Box<RawValue>>",
+        "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<Box<RawValue>>",
+        "type_info": "Jsonb"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Uuid",
+        "Text"
+      ]
+    },
+    "nullable": [
+      true,
+      true,
+      true,
+      true,
+      true,
+      true,
+      true,
+      true
+    ]
+  },
+  "hash": "f3d8c6804a4d3ebf69704ebe7956c38ca756fdc998cb1fe3e79f54269a914b49"
+}

backend/windmill-api/src/jobs.rs

@@ -8,7 +8,6 @@
 
 use axum::body::Body;
 use axum::http::HeaderValue;
-use futures::TryFutureExt;
 use http::{HeaderMap, HeaderName};
 use itertools::Itertools;
 use quick_cache::sync::Cache;
@@ -795,38 +794,25 @@ impl<'a> GetQuery<'a> {
     async fn resolve_raw_values<T>(
         &self,
         db: &DB,
-        id: Uuid,
         kind: JobKind,
-        hash: Option<ScriptHash>,
+        runnable_id: Option<i64>,
         job: &mut JobExtended<T>,
     ) {
-        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,
+            flow => cache::job::fetch_flow(db, kind, runnable_id, 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),
+            (code, lock) => cache::job::fetch_script(db, kind, runnable_id, code, lock)
                 .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(
@@ -849,7 +835,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 {
@@ -881,7 +867,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 {

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},
 };
 
@@ -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_as_queue.job_kind AS \"job_kind!: JobKind\",
-            v2_as_queue.script_hash AS \"script_hash: ScriptHash\",
+            v2_as_queue.script_hash,
             v2_as_queue.raw_flow AS \"raw_flow: sqlx::types::Json<Box<RawValue>>\",
             v2_as_completed_job.parent_job AS \"parent_job: Uuid\",
             v2_as_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;

backend/windmill-common/src/cache.rs

@@ -302,6 +302,26 @@ pub enum RawData {
     Script(Arc<ScriptData>),
 }
 
+impl RawData {
+    pub fn from_raw(
+        raw_code: Option<String>,
+        raw_lock: Option<String>,
+        raw_flow: Option<Json<Box<RawValue>>>,
+    ) -> error::Result<Option<Self>> {
+        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<ScriptLang>,
@@ -353,7 +373,6 @@ pub fn clear() {
     flow::clear();
     script::clear();
     app::clear();
-    job::clear();
 }
 
 pub mod flow {
@@ -543,16 +562,10 @@ pub mod script {
                 meta: Some(ScriptMetadata {
                     language: r.language,
                     envs: r.envs,
-                    codebase: if let Some(use_tar) = r.use_tar {
+                    codebase: r.use_tar.map(|use_tar| {
                         let sh = hash.to_string();
-                        if use_tar {
-                            Some(format!("{sh}.tar"))
-                        } else {
-                            Some(sh)
-                        }
-                    } else {
-                        None
-                    },
+                        use_tar.then(|| format!("{sh}.tar")).unwrap_or(sh)
+                    }),
                 }),
             })
         });
@@ -610,114 +623,21 @@ 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<Uuid, RawData> = 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<Json<Box<RawValue>>>,
-    ) -> impl Future<Output = error::Result<Arc<FlowData>>> + '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::internal_err(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<String>,
-        raw_code: Option<String>,
-    ) -> impl Future<Output = error::Result<Arc<ScriptData>>> + '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::internal_err(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<String>,
-        raw_code: Option<String>,
-        raw_flow: Option<Json<Box<RawValue>>>,
-    ) -> impl Future<Output = error::Result<RawData>> + '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<Box<RawValue>>\" \
-                    FROM v2_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
-        }
-    }
-
     #[track_caller]
     pub fn fetch_script<'c>(
         e: impl PgExecutor<'c>,
         kind: JobKind,
-        hash: Option<ScriptHash>,
+        runnable_id: Option<i64>,
+        raw_code: Option<String>,
+        raw_lock: Option<String>,
     ) -> impl Future<Output = error::Result<Arc<ScriptData>>> {
         use JobKind::*;
         let loc = Location::caller();
         async move {
-            match (kind, hash.map(|ScriptHash(id)| id)) {
+            if let (Some(code), lock) = (raw_code, raw_lock) {
+                return Ok(Arc::new(ScriptData { code, lock }));
+            }
+            match (kind, runnable_id) {
                 (FlowScript, Some(id)) => flow::fetch_script(e, FlowNodeId(id)).await,
                 (Script | Dependencies, Some(hash)) => script::fetch(e, ScriptHash(hash))
                     .await
@@ -736,16 +656,20 @@ pub mod job {
     pub fn fetch_flow<'c>(
         e: impl PgExecutor<'c> + Copy,
         kind: JobKind,
-        hash: Option<ScriptHash>,
+        runnable_id: Option<i64>,
+        raw_flow: Option<Json<Box<RawValue>>>,
     ) -> impl Future<Output = error::Result<Arc<FlowData>>> {
         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::internal_err(format!(

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 as FlowIterator, JobResult, RestartedFrom, RetryStatus, MAX_RETRY_ATTEMPTS,
+        BranchAllStatus, FlowStatus, FlowStatusModule, FlowStatusModuleWParent, Iterator as FlowIterator,
+        JobResult, RestartedFrom, RetryStatus, MAX_RETRY_ATTEMPTS,
         MAX_RETRY_INTERVAL,
     },
     flows::{
@@ -631,7 +631,7 @@ pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
             )
             .execute(&mut *tx)
             .await
-            .map_err(|e| Error::InternalErr(format!("Could not update job labels: {e:#}")))?;
+            .map_err(|e| Error::internal_err(format!("Could not update job labels: {e:#}")))?;
         }
 
         if !queued_job.is_flow_step {
@@ -919,7 +919,9 @@ pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
     #[cfg(feature = "enterprise")]
     if !success {
         async fn has_failure_module(db: &Pool<Postgres>, 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!(
@@ -3229,41 +3231,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<Box<RawValue>>\"
+                            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)
@@ -3443,45 +3424,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<Box<RawValue>>\"
+                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::internal_err(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();
@@ -3490,19 +3452,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,
@@ -3933,50 +3895,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<Postgres>,
+fn restarted_flows_resolution(
     workspace_id: &str,
-    completed_flow_id: Uuid,
-    restart_step_id: &str,
-    branch_or_iteration_n: Option<usize>,
-) -> Result<
-    (
-        Option<i64>,
-        Option<String>,
-        Arc<FlowData>,
-        i32,
-        Vec<FlowStatusModule>,
-        HashMap<String, serde_json::Value>,
-        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<Box<RawValue>>\",
-            raw_flow AS \"raw_flow: Json<Box<RawValue>>\"
-        FROM v2_as_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::internal_err(format!(
-            "completed job not found for UUID {} in workspace {}: {}",
-            completed_flow_id, workspace_id, err
-        ))
-    })?;
+    flow_value: &FlowValue,
+    flow_status: Option<Json<Box<RawValue>>>,
+    restart_from: RestartedFrom,
+) -> error::Result<FlowStatus> {
+    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::<FlowStatus>(v.get()).ok())
         .ok_or(Error::internal_err(format!(
             "Unable to parse flow status for job {} in workspace {}",
@@ -3987,20 +3918,21 @@ async fn restarted_flows_resolution(
     let mut dependent_module = false;
     let mut truncated_modules: Vec<FlowStatusModule> = 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();
@@ -4029,7 +3961,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),
@@ -4067,7 +3999,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(FlowIterator {
                                 index: branch_or_iteration_n - 1, // same deal as above, this refers to the last finished job
@@ -4093,7 +4025,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;
@@ -4114,13 +4046,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)
 }

backend/windmill-worker/src/result_processor.rs

@@ -451,7 +451,7 @@ pub async fn process_completed_job(
             .execute(db)
             .await
             .map_err(|e| {
-                Error::InternalErr(format!(
+                Error::internal_err(format!(
                     "error while deleting args of preprocessing step: {e:#}"
                 ))
             })?;

backend/windmill-worker/src/worker.rs

@@ -2045,22 +2045,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
     };
@@ -2102,10 +2089,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,
@@ -2164,7 +2152,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,
@@ -2180,7 +2168,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,
@@ -2216,13 +2204,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,

backend/windmill-worker/src/worker_flow.rs

@@ -20,7 +20,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;
@@ -42,7 +41,6 @@ use windmill_common::jobs::{
     script_hash_to_tag_and_limits, script_path_to_payload, 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;
@@ -218,7 +216,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<Box<RawValue>>\",
                 raw_flow AS \"raw_flow: Json<Box<RawValue>>\"
             FROM v2_as_queue WHERE id = $1 AND workspace_id = $2 LIMIT 1",
@@ -245,9 +243,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());
@@ -1047,7 +1043,7 @@ pub async fn update_flow_status_after_job_completion_internal(
                 .execute(db)
                 .await
                 .map_err(|e| {
-                    Error::InternalErr(format!("error while cleaning up completed job: {e:#}"))
+                    Error::internal_err(format!("error while cleaning up completed job: {e:#}"))
                 })?;
                 sqlx::query!(
                     "UPDATE v2_job_completed SET result = '{}'::jsonb WHERE id = ANY($1)",