From 75f2e39418ee9a3b20c6c1d09eca5e9288ea636c Mon Sep 17 00:00:00 2001
From: Abel Lucas <uaelucas@gmail.com>
Date: Thu, 6 Feb 2025 15:27:41 +0100
Subject: [PATCH] backend: make `pull` atomic again

---
 backend/windmill-queue/src/jobs.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs
index 48f66fcf5b..28f9f22630 100644
--- a/backend/windmill-queue/src/jobs.rs
+++ b/backend/windmill-queue/src/jobs.rs
@@ -2060,6 +2060,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
     db: &Pool<Postgres>,
     suspend_first: bool,
 ) -> windmill_common::error::Result<(Option<PulledJob>, bool)> {
+    let mut tx = db.begin().await?;
     let job_and_suspended: (Option<PulledJob>, bool) = {
         /* Jobs can be started if they:
          * - haven't been started before,
@@ -2078,7 +2079,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
         let r = if suspend_first {
             // tracing::info!("Pulling job with query: {}", query);
             sqlx::query_as::<_, PulledJob>(&query)
-                .fetch_optional(db)
+                .fetch_optional(&mut *tx)
                 .await?
         } else {
             None
@@ -2098,7 +2099,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
             for query in queries.iter() {
                 // tracing::info!("Pulling job with query: {}", query);
                 let r = sqlx::query_as::<_, PulledJob>(query)
-                    .fetch_optional(db)
+                    .fetch_optional(&mut *tx)
                     .await?;
 
                 if let Some(pulled_job) = r {
@@ -2115,6 +2116,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
             (r, true)
         }
     };
+    tx.commit().await?;
     Ok(job_and_suspended)
 }