From dead8df080d7f731660688f73e4492dfc1a33255 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 24 Aug 2026 17:30:49 +0000 Subject: [PATCH] fix: make Pool::begin cancel-safe across the backend Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Qzqmh52NU8fB9RBQNNkJGt --- backend/src/main.rs | 3 +- backend/src/monitor.rs | 15 ++-- backend/windmill-api-configs/src/lib.rs | 5 +- backend/windmill-api-flows/src/flows.rs | 3 +- backend/windmill-api-groups/src/groups.rs | 19 ++--- .../src/concurrency_groups.rs | 3 +- backend/windmill-api-jobs/src/execution.rs | 7 +- backend/windmill-api-schedule/src/lib.rs | 5 +- backend/windmill-api-scripts/src/scripts.rs | 9 ++- .../windmill-api-settings/src/log_cleanup.rs | 3 +- backend/windmill-api-users/src/users.rs | 49 ++++++------ .../src/datatable_migrations.rs | 11 +-- .../src/deployment_requests.rs | 9 ++- .../windmill-api-workspaces/src/workspaces.rs | 79 ++++++++++--------- .../src/workspaces_extra.rs | 9 ++- backend/windmill-api/src/ai.rs | 9 ++- backend/windmill-api/src/ai_evals/results.rs | 3 +- backend/windmill-api/src/ai_evals/run.rs | 3 +- backend/windmill-api/src/ai_skills.rs | 5 +- backend/windmill-api/src/apps.rs | 5 +- backend/windmill-api/src/jobs.rs | 17 ++-- backend/windmill-api/src/offboarding.rs | 5 +- backend/windmill-api/src/users.rs | 5 +- .../windmill-api/src/workspace_shared_ui.rs | 3 +- backend/windmill-api/src/workspaces.rs | 3 +- backend/windmill-common/src/bench.rs | 3 +- backend/windmill-common/src/db.rs | 42 +++++++++- backend/windmill-common/src/dbt_manifest.rs | 3 +- backend/windmill-common/src/scripts.rs | 3 +- backend/windmill-common/src/usernames.rs | 3 +- backend/windmill-common/src/utils.rs | 3 +- .../tests/begin_cancel_safe.rs | 60 ++++++++++++++ backend/windmill-dep-map/src/lib.rs | 3 +- .../src/scoped_dependency_map.rs | 7 +- .../src/workspace_dependencies.rs | 3 +- .../src/google/external.rs | 3 +- .../windmill-native-triggers/src/rename.rs | 5 +- .../src/workspace_integrations.rs | 3 +- backend/windmill-queue/src/jobs.rs | 9 ++- backend/windmill-queue/src/schedule.rs | 3 +- backend/windmill-store/src/resources.rs | 3 +- backend/windmill-trigger/src/listener.rs | 3 +- backend/windmill-worker/src/ai/tools.rs | 3 +- backend/windmill-worker/src/ai/utils.rs | 3 +- backend/windmill-worker/src/ai_executor.rs | 3 +- backend/windmill-worker/src/bun_executor.rs | 9 ++- backend/windmill-worker/src/dbt_executor.rs | 7 +- .../windmill-worker/src/duckdb_executor.rs | 3 +- .../windmill-worker/src/result_processor.rs | 3 +- backend/windmill-worker/src/worker_flow.rs | 15 ++-- .../windmill-worker/src/worker_lockfiles.rs | 13 +-- 51 files changed, 325 insertions(+), 180 deletions(-) create mode 100644 backend/windmill-common/tests/begin_cancel_safe.rs diff --git a/backend/src/main.rs b/backend/src/main.rs index 1819146cbe..17cee0f889 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -5,6 +5,7 @@ * Please see the included NOTICE for copyright information and * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use anyhow::Context; use monitor::{ load_base_url, load_otel, reload_critical_alerts_on_db_oversize, @@ -2401,7 +2402,7 @@ async fn coordinate_restart_delay( let now = chrono::Utc::now(); - let mut tx = db.begin().await.context("begin restart coordination tx")?; + let mut tx = db.begin_cancel_safe().await.context("begin restart coordination tx")?; // Serialize access across all instances sqlx::query("SELECT pg_advisory_xact_lock($1)") diff --git a/backend/src/monitor.rs b/backend/src/monitor.rs index 94dcbf0d46..51adff243d 100644 --- a/backend/src/monitor.rs +++ b/backend/src/monitor.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; #[cfg(feature = "oauth2")] use std::collections::HashMap; use std::{ @@ -1837,7 +1838,7 @@ async fn cleanup_scheduled_job_deletions(db: &Pool) { let mut total_deleted = 0u64; for batch_num in 0..MAX_BATCHES { - let mut tx = match db.begin().await { + let mut tx = match db.begin_cancel_safe().await { Ok(tx) => tx, Err(e) => { tracing::error!("Error starting transaction for scheduled job deletion: {e:?}"); @@ -1976,7 +1977,7 @@ async fn delete_expired_jobs_batch( only_workspace: Option<&str>, exclude_workspaces: Option<&[String]>, ) -> error::Result<(usize, Option>)> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Fetch active ROOT job IDs that started before the retention period. We only care about // these because their child jobs could be old enough to be deletion candidates. @@ -4191,7 +4192,7 @@ async fn reconcile_unarmed_schedules(db: &Pool) { // rollback a dropped `Transaction` performs — so cancellation can't strand it. // The tx is held open only to own the lock; the scan and re-arm run on separate // pool connections. - let mut lock_tx = match db.begin().await { + let mut lock_tx = match db.begin_cancel_safe().await { Ok(tx) => tx, Err(e) => { tracing::error!("schedule reconcile: failed to begin lock tx: {e:#}"); @@ -5043,7 +5044,7 @@ async fn cancel_stale_job( workspace_id: String, scheduled_for: DateTime, ) -> error::Result<()> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; tracing::error!( "Stale job detected: {} in workspace {} with tag {} (scheduled for: {}) . Cancelling it.", id, @@ -5410,7 +5411,7 @@ async fn force_complete_zombie_job( "name": "ExecutionErr", }); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "INSERT INTO v2_job_completed @@ -5663,7 +5664,7 @@ async fn handle_zombie_flows(db: &DB) -> error::Result<()> { tracing::error!(error_message); report_critical_error(error_message, db.clone(), Some(&flow.workspace_id), None).await; // if the flow hasn't started and is a zombie, we can simply restart it - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let concurrency_key = sqlx::query_scalar!("SELECT key FROM concurrency_key WHERE job_id = $1", flow.id) @@ -6126,7 +6127,7 @@ async fn cancel_zombie_flow_job( workspace_id: &str, message: String, ) -> Result<(), error::Error> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; tracing::error!( "zombie flow detected: {} in workspace {}. Cancelling it.", id, diff --git a/backend/windmill-api-configs/src/lib.rs b/backend/windmill-api-configs/src/lib.rs index 7ad99996d5..64c9371693 100644 --- a/backend/windmill-api-configs/src/lib.rs +++ b/backend/windmill-api-configs/src/lib.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use axum::{ extract::{Extension, Path, Query}, routing::{get, post}, @@ -187,7 +188,7 @@ async fn update_config( } } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "INSERT INTO config (name, config) VALUES ($1, $2) ON CONFLICT (name) DO UPDATE SET config = EXCLUDED.config", &name, @@ -217,7 +218,7 @@ async fn delete_config( ) -> error::Result { require_devops_role(&db, &authed).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let deleted = sqlx::query!("DELETE FROM config WHERE name = $1 RETURNING name", name) .fetch_all(&mut *tx) diff --git a/backend/windmill-api-flows/src/flows.rs b/backend/windmill-api-flows/src/flows.rs index 6353cd429a..4113529245 100644 --- a/backend/windmill-api-flows/src/flows.rs +++ b/backend/windmill-api-flows/src/flows.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use std::collections::HashMap; use axum::response::IntoResponse; @@ -1589,7 +1590,7 @@ async fn get_deployment_status( Path((w_id, path)): Path<(String, StripPath)>, ) -> JsonResult { let path = path.to_path(); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let status_o = sqlx::query!( "SELECT f.lock_error_logs, dm.job_id FROM flow f diff --git a/backend/windmill-api-groups/src/groups.rs b/backend/windmill-api-groups/src/groups.rs index a15b9823d3..a6a8f0c13f 100644 --- a/backend/windmill-api-groups/src/groups.rs +++ b/backend/windmill-api-groups/src/groups.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use windmill_api_auth::{require_super_admin, ApiAuthed}; use windmill_common::DB; @@ -300,7 +301,7 @@ async fn create_igroup( use uuid::Uuid; require_super_admin(&db, &authed).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let normalized_name = convert_name(&ng.name); @@ -465,7 +466,7 @@ async fn update_igroup( Json(igroup_update): Json, ) -> Result { require_super_admin(&db, &authed).await?; - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; let exists_opt = sqlx::query("SELECT 1 FROM instance_group WHERE name = $1") .bind(name.clone()) @@ -657,7 +658,7 @@ async fn delete_igroup( Path(name): Path, ) -> Result { require_super_admin(&db, &authed).await?; - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; // FOR UPDATE: the group row is the group-level mutex, taken before the workspace // advisory locks (see reconcile_workspace_instance_groups). @@ -972,7 +973,7 @@ async fn add_user_igroup( ) -> Result { require_super_admin(&db, &authed).await?; - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; // FOR UPDATE: the group row is the group-level mutex, taken before the workspace // advisory locks (see reconcile_workspace_instance_groups). @@ -1048,7 +1049,7 @@ struct WorkspaceInfo { role: String, } async fn list_igroups(Extension(db): Extension) -> JsonResult> { - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; let groups = sqlx::query_as!( IGroup, @@ -1064,7 +1065,7 @@ async fn list_igroups(Extension(db): Extension) -> JsonResult> { async fn list_igroups_with_workspaces( Extension(db): Extension, ) -> JsonResult> { - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; // Get all instance groups with their emails first let groups = sqlx::query_as!( @@ -1190,7 +1191,7 @@ async fn remove_user_igroup( Json(Email { email }): Json, ) -> Result { require_super_admin(&db, &authed).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // FOR UPDATE: the group row is the group-level mutex, taken before the workspace // advisory locks (see reconcile_workspace_instance_groups). @@ -1331,7 +1332,7 @@ async fn export_igroups( Extension(db): Extension, ) -> JsonResult> { require_super_admin(&db, &authed).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let igroups = sqlx::query_as!( ExportedIGroup, "SELECT name, summary, array_remove(array_agg(email_to_igroup.email), null) as emails, id, scim_display_name, external_id, instance_role FROM email_to_igroup RIGHT JOIN instance_group ON instance_group.name = email_to_igroup.igroup GROUP BY name", @@ -1367,7 +1368,7 @@ async fn overwrite_igroups( Json(igroups): Json>, ) -> Result { require_super_admin(&db, &authed).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // The import replaces the whole group catalog, so the whole-table lock is its // group-mutex phase, taken first like every path's group locks (see diff --git a/backend/windmill-api-jobs/src/concurrency_groups.rs b/backend/windmill-api-jobs/src/concurrency_groups.rs index 4040d41bd2..9e6f37c2ca 100644 --- a/backend/windmill-api-jobs/src/concurrency_groups.rs +++ b/backend/windmill-api-jobs/src/concurrency_groups.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use windmill_api_auth::{check_scopes, is_instance_admin, require_instance_admin, ApiAuthed}; use windmill_common::{ db::{UserDB, DB}, @@ -75,7 +76,7 @@ async fn prune_concurrency_group( "Only administrators can delete concurrency groups".to_string(), )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let concurrency_group = sqlx::query_as::<_, (String, i64)>( "SELECT concurrency_id, (select COUNT(*) from jsonb_object_keys(job_uuids)) as n_job_uuids FROM concurrency_counter WHERE concurrency_id = $1 FOR UPDATE", diff --git a/backend/windmill-api-jobs/src/execution.rs b/backend/windmill-api-jobs/src/execution.rs index fbfd656cdf..22b568266f 100644 --- a/backend/windmill-api-jobs/src/execution.rs +++ b/backend/windmill-api-jobs/src/execution.rs @@ -8,6 +8,7 @@ //! Job execution primitives: push, cancel, tag/license checks, wait-for-result. +use windmill_common::db::BeginCancelSafe; use axum::{ response::{IntoResponse, Response}, Json, @@ -91,7 +92,7 @@ pub async fn cancel_jobs( ) -> error::JsonResult> { let mut uuids = vec![]; tracing::info!("Cancelling jobs: {:?}", jobs); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let trivial_jobs = sqlx::query!("INSERT INTO v2_job_completed AS cj ( workspace_id , id @@ -134,7 +135,7 @@ pub async fn cancel_jobs( continue; } match tokio::time::timeout(tokio::time::Duration::from_secs(5), async move { - let tx = db.begin().await?; + let tx = db.begin_cancel_safe().await?; let (tx, _) = cancel_job( username, None, @@ -196,7 +197,7 @@ impl Drop for Guard { tracing::info!("http connection broke, marking job {id} as canceled"); tokio::spawn(async move { let cancel_f = async { - let tx = db.begin().await?; + let tx = db.begin_cancel_safe().await?; let (tx, _) = cancel_job( &username, Some("http connection broke".to_string()), diff --git a/backend/windmill-api-schedule/src/lib.rs b/backend/windmill-api-schedule/src/lib.rs index 6ae902e984..19ddb4f2aa 100644 --- a/backend/windmill-api-schedule/src/lib.rs +++ b/backend/windmill-api-schedule/src/lib.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use axum::{ extract::{Extension, Path, Query}, routing::{delete, get, post}, @@ -1087,7 +1088,7 @@ async fn exists_schedule( Extension(db): Extension, Path((w_id, path)): Path<(String, StripPath)>, ) -> JsonResult { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let res = windmill_queue::schedule::exists_schedule(&mut tx, w_id, path).await?; tx.commit().await?; Ok(Json(res)) @@ -1482,7 +1483,7 @@ async fn set_default_error_handler( // connections a concurrent edit could interleave, leaving the // id-ordered drawer showing the wrong latest change, and a failed // insert would leave the schedules rewritten with nothing recording it. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let updated_schedules: Vec; match payload.handler_type { HandlerType::Error => { diff --git a/backend/windmill-api-scripts/src/scripts.rs b/backend/windmill-api-scripts/src/scripts.rs index 79c6c99b9b..7909f4ad75 100644 --- a/backend/windmill-api-scripts/src/scripts.rs +++ b/backend/windmill-api-scripts/src/scripts.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use axum::extract::Multipart; use windmill_api_auth::{ auth::{list_tokens_internal, AuthCache, TruncatedTokenWithEmail}, @@ -3447,7 +3448,7 @@ async fn get_script_by_hash( let mut tx = if query_auth.authed.is_some_and(|x| x) { user_db.begin(&authed).await? } else { - db.begin().await? + db.begin_cancel_safe().await? }; let r = get_script_by_hash_internal( &mut tx, @@ -3476,7 +3477,7 @@ async fn raw_script_by_hash( Extension(db): Extension, Path((w_id, hash_str)): Path<(String, String)>, ) -> Result { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let hash = ScriptHash(to_i64(hash_str.strip_suffix(".ts").ok_or_else(|| { Error::BadRequest("Raw script path must end with .ts".to_string()) })?)?); @@ -3496,7 +3497,7 @@ async fn get_deployment_status( Extension(db): Extension, Path((w_id, hash)): Path<(String, ScriptHash)>, ) -> JsonResult { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let status_o = sqlx::query!( "SELECT s.lock, s.lock_error_logs, dm.job_id FROM script s @@ -3961,7 +3962,7 @@ async fn delete_scripts_bulk( return Err(Error::PermissionDenied(msg)); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Capture scripts for trashbin per path before bulk delete for path in &request.paths { diff --git a/backend/windmill-api-settings/src/log_cleanup.rs b/backend/windmill-api-settings/src/log_cleanup.rs index ecbc38fcb7..9d0657195d 100644 --- a/backend/windmill-api-settings/src/log_cleanup.rs +++ b/backend/windmill-api-settings/src/log_cleanup.rs @@ -17,6 +17,7 @@ * the status endpoint — not just the one that happened to receive the POST. */ +use windmill_common::db::BeginCancelSafe; use std::sync::Arc; use chrono::{DateTime, Utc}; @@ -460,7 +461,7 @@ async fn delete_expired_jobs_batch( only_workspace: Option<&str>, exclude_workspaces: Option<&[String]>, ) -> error::Result<(usize, Vec, Option>)> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let active_root_job_ids: Vec = sqlx::query_scalar!( "SELECT q.id FROM v2_job_queue q diff --git a/backend/windmill-api-users/src/users.rs b/backend/windmill-api-users/src/users.rs index 8ea3bd0995..834610a9a0 100644 --- a/backend/windmill-api-users/src/users.rs +++ b/backend/windmill-api-users/src/users.rs @@ -8,6 +8,7 @@ #![allow(non_snake_case)] +use windmill_common::db::BeginCancelSafe; use sqlx::{Postgres, Transaction}; use std::sync::atomic::AtomicBool; @@ -611,7 +612,7 @@ async fn list_invites( authed: ApiAuthed, Extension(db): Extension, ) -> JsonResult> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let rows = sqlx::query_as!( WorkspaceInvite, "SELECT @@ -652,7 +653,7 @@ async fn logout( cookie.set_domain(COOKIE_DOMAIN.clone().unwrap()); } cookies.remove(cookie); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let t_hash = windmill_common::auth::hash_token(&token); let t_prefix = token.get(..TOKEN_PREFIX_LEN).unwrap_or(&token); @@ -976,7 +977,7 @@ async fn whois( // Json(nu): Json, // ) -> Result<(StatusCode, String)> { -// let mut tx = db.begin().await?; +// let mut tx = db.begin_cancel_safe().await?; // require_super_admin(&mut *tx, email).await?; // sqlx::query!( @@ -1002,7 +1003,7 @@ async fn decline_invite( Extension(db): Extension, Json(nu): Json, ) -> Result<(StatusCode, String)> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let is_admin = sqlx::query_scalar!( "DELETE FROM workspace_invite WHERE workspace_id = $1 AND email = $2 RETURNING is_admin", @@ -1048,7 +1049,7 @@ async fn accept_invite( Extension(db): Extension, Json(nu): Json, ) -> Result<(StatusCode, String)> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let r = sqlx::query!( "DELETE FROM workspace_invite WHERE workspace_id = $1 AND email = $2 RETURNING is_admin, operator", @@ -1235,7 +1236,7 @@ async fn join_workspace<'c>( async fn leave_instance(Extension(db): Extension, authed: ApiAuthed) -> Result { forbid_job_token_account_destruction(&authed)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!("DELETE FROM password WHERE email = $1", &authed.email) .execute(&mut *tx) .await?; @@ -1282,7 +1283,7 @@ async fn update_workspace_user( Path((w_id, username_to_update)): Path<(String, String)>, Json(eu): Json, ) -> Result { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; require_admin(authed.is_admin, &authed.username)?; @@ -1328,7 +1329,7 @@ async fn convert_user_to_group( Path((w_id, username_to_convert)): Path<(String, String)>, ) -> Result { require_admin(authed.is_admin, &authed.username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Get user email and current status let user_info = sqlx::query!( @@ -1477,7 +1478,7 @@ async fn update_user( ) -> Result { require_super_admin(&db, &authed).await?; forbid_superadmin_job_token(&db, &authed.email, job_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let mut new_super_admin: Option = None; if let Some(sa) = eu.is_super_admin { @@ -1653,7 +1654,7 @@ async fn delete_user( ) -> Result { require_super_admin(&db, &authed).await?; forbid_superadmin_job_token(&db, &authed.email, job_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!("DELETE FROM token WHERE email = $1", &email_to_delete) .execute(&mut *tx) @@ -1773,7 +1774,7 @@ async fn change_user_email( } } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // FOR UPDATE serializes concurrent moves of *this* account. Two moves of different accounts // onto the same destination are stopped by the `password` primary key instead, which is why the @@ -2548,7 +2549,7 @@ async fn delete_workspace_user( Extension(db): Extension, Path((w_id, username_to_delete)): Path<(String, String)>, ) -> Result { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Locked so that the authorization below and the delete it guards see the same row. let target = sqlx::query!( @@ -2609,7 +2610,7 @@ async fn set_login_type( ) -> Result { require_super_admin(&db, &authed).await?; forbid_superadmin_job_token(&db, &authed.email, job_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE password SET login_type = $1 WHERE email = $2", @@ -2664,7 +2665,7 @@ async fn login( windmill_common::login_rate_limit::check_and_increment_login_attempt(&headers, &email)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let audit_author = AuditAuthor { email: email.clone(), username: email.clone(), @@ -2774,7 +2775,7 @@ async fn refresh_token( } } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let super_admin = sqlx::query_scalar!( "SELECT super_admin FROM password WHERE email = $1 AND disabled = false", @@ -2925,7 +2926,7 @@ async fn create_token( windmill_api_auth::ensure_scopes_within_caller(&authed, token_config.scopes.as_deref())?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let token = create_token_internal(&mut *tx, &db, &authed, token_config).await?; @@ -2981,7 +2982,7 @@ async fn impersonate( .fetch_optional(&db) .await? .unwrap_or(false); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "INSERT INTO token @@ -3105,7 +3106,7 @@ async fn delete_token( Path(token_prefix): Path, ) -> Result { forbid_job_token_account_destruction(&authed)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let tokens_deleted: Vec = sqlx::query_scalar( "DELETE FROM token @@ -3156,7 +3157,7 @@ async fn update_token_scopes( forbid_elevated_job_token(&db, &authed.email, authed.job_id).await?; windmill_api_auth::ensure_scopes_within_caller(&authed, req.scopes.as_deref())?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let updated: Option = sqlx::query_scalar!( "UPDATE token SET scopes = $1 @@ -3230,7 +3231,7 @@ async fn update_token_label( ))); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Only user-created tokens may be relabeled — system tokens carry the // load-bearing labels described above. This SQL mirrors the canonical @@ -3283,7 +3284,7 @@ async fn leave_workspace( authed: ApiAuthed, ) -> Result { forbid_job_token_account_destruction(&authed)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "DELETE FROM usr WHERE workspace_id = $1 AND username = $2", &w_id, @@ -3429,7 +3430,7 @@ async fn get_instance_username_info( Extension(db): Extension, ) -> JsonResult { require_super_admin(&db, &authed).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let instance_username = match sqlx::query_scalar!( "SELECT username FROM password WHERE email = $1", &user_email @@ -3500,7 +3501,7 @@ async fn export_global_users( ) -> JsonResult> { require_super_admin(&db, &authed).await?; forbid_superadmin_job_token(&db, &authed.email, job_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let users = sqlx::query_as!( ExportedGlobalUser, "SELECT email, password_hash, login_type, super_admin, verified, name, company, first_time_user, username FROM password" @@ -3540,7 +3541,7 @@ async fn overwrite_global_users( ) -> Result { require_super_admin(&db, &authed).await?; forbid_superadmin_job_token(&db, &authed.email, job_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!("DELETE FROM password") .execute(&mut *tx) .await?; diff --git a/backend/windmill-api-workspaces/src/datatable_migrations.rs b/backend/windmill-api-workspaces/src/datatable_migrations.rs index 14c6d49132..f5a12bdeb1 100644 --- a/backend/windmill-api-workspaces/src/datatable_migrations.rs +++ b/backend/windmill-api-workspaces/src/datatable_migrations.rs @@ -10,6 +10,7 @@ //! management, and the workspace-merge diff helper. Split out of `workspaces.rs` //! to keep that file focused on core workspace configuration. +use windmill_common::db::BeginCancelSafe; use crate::workspaces::{pg_dump_database, ItemComparison}; use axum::{ @@ -834,7 +835,7 @@ async fn disable_datatable_migrations( ) -> Result { require_datatable_migrations_manager(&db, &authed).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let updated = sqlx::query_scalar!( "UPDATE workspace_settings \ @@ -1098,7 +1099,7 @@ async fn create_datatable_migration( validate_migration_name(&payload.name)?; ensure_datatable_migrations_enabled(&db, &w_id, &datatable_name).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let timestamp = insert_datatable_migration_def( &mut tx, &w_id, @@ -1390,7 +1391,7 @@ async fn generate_initial_datatable_migration( // than a `_wm_migrations` version with no definition that the UI can't // clear). The narrow window where it briefly shows "not run" is benign: // running it would just no-op/fail harmlessly against the existing schema. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let timestamp = insert_datatable_migration_def(&mut tx, &w_id, &datatable_name, "initial", &code_up, None) .await?; @@ -1861,7 +1862,7 @@ mod tests { seed_migration(&pool, &w_id, "sa", 5000, "sa_mig").await; seed_migration(&pool, &w_id, "sb", 5000, "sb_mig").await; - let mut tx = pool.begin().await.unwrap(); + let mut tx = pool.begin_cancel_safe().await.unwrap(); let mut changed = cascade_datatable_migration_renames_and_deletes( &pool, &mut tx, @@ -1919,7 +1920,7 @@ mod tests { seed_migration(&pool, &src, "customers", 2, "add_index").await; seed_migration(&pool, &src, "orders", 3, "create_orders").await; - let mut tx = pool.begin().await.unwrap(); + let mut tx = pool.begin_cancel_safe().await.unwrap(); clone_datatable_migrations(&mut tx, &src, &dst) .await .unwrap(); diff --git a/backend/windmill-api-workspaces/src/deployment_requests.rs b/backend/windmill-api-workspaces/src/deployment_requests.rs index a279f82d6c..27a60a3724 100644 --- a/backend/windmill-api-workspaces/src/deployment_requests.rs +++ b/backend/windmill-api-workspaces/src/deployment_requests.rs @@ -10,6 +10,7 @@ //! Email dispatch happens via `send_email_if_possible`, which is a no-op on //! OSS builds. +use windmill_common::db::BeginCancelSafe; use axum::{ extract::{Extension, Path}, http::StatusCode, @@ -283,7 +284,7 @@ async fn create_deployment_request( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let request_id_row = sqlx::query!( r#" @@ -399,7 +400,7 @@ async fn cancel_deployment_request( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let rows_affected = sqlx::query!( "UPDATE workspace_fork_deployment_request SET closed_at = now(), closed_reason = 'cancelled' WHERE id = $1 AND closed_at IS NULL", id, @@ -499,7 +500,7 @@ async fn close_deployment_request_merged( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let rows_affected = sqlx::query!( "UPDATE workspace_fork_deployment_request SET closed_at = now(), closed_reason = 'merged' WHERE id = $1 AND closed_at IS NULL", id, @@ -641,7 +642,7 @@ async fn create_deployment_request_comment( None }; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let row = sqlx::query!( r#" INSERT INTO workspace_fork_deployment_request_comment diff --git a/backend/windmill-api-workspaces/src/workspaces.rs b/backend/windmill-api-workspaces/src/workspaces.rs index e04e1c1ab9..006506836b 100644 --- a/backend/windmill-api-workspaces/src/workspaces.rs +++ b/backend/windmill-api-workspaces/src/workspaces.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use windmill_api_auth::{ build_scope_path_predicate, check_scopes, require_devops_role, require_instance_admin, require_is_writer, require_super_admin, ApiAuthed, @@ -1466,7 +1467,7 @@ async fn edit_slack_command( Json(es): Json, ) -> Result { require_admin(is_admin, &username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; if es.slack_command_script.is_some() { let exists_slack_command_with_team_id = sqlx::query_scalar!( @@ -1614,7 +1615,7 @@ async fn set_slack_oauth_config( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE workspace_settings @@ -1662,7 +1663,7 @@ async fn delete_slack_oauth_config( ) -> Result { require_admin(authed.is_admin, &authed.username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE workspace_settings @@ -1798,7 +1799,7 @@ async fn edit_webhook( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; if let Some(webhook) = &ew.webhook { sqlx::query!( @@ -1909,7 +1910,7 @@ async fn edit_dbt_warehouses( .map(|m| serde_json::to_value(m)) .transpose() .map_err(|e| Error::internal_err(format!("serializing the dbt warehouses: {e}")))?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; audit_log( &mut *tx, &authed, @@ -1941,7 +1942,7 @@ async fn edit_large_file_storage_config( ) -> Result { require_admin(is_admin, &username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let args_for_audit = format!("{:?}", new_config.large_file_storage); audit_log( @@ -3064,7 +3065,7 @@ async fn edit_ducklake_config( } } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let args_for_audit = format!("{:?}", new_config.settings); audit_log( @@ -3150,7 +3151,7 @@ async fn edit_datatable_config( require_admin(is_admin, &username)?; let is_superadmin = require_super_admin(&db, &authed).await.is_ok(); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let old_datatables: HashMap = serde_json::from_value( sqlx::query_scalar!( @@ -3608,7 +3609,7 @@ async fn edit_git_sync_config( require_admin(is_admin, &username)?; check_git_sync_access(&db, &w_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let args_for_audit = format!("{:?}", new_config.git_sync_settings); audit_log( @@ -3916,7 +3917,7 @@ async fn edit_git_sync_repository( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // First, get the current git sync settings. `FOR UPDATE` because this // read-modify-writes the whole `git_sync` blob: the webhook reconciler writes @@ -4096,7 +4097,7 @@ async fn delete_git_sync_repository( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // First, get the current git sync settings. `FOR UPDATE` because this // read-modify-writes the whole `git_sync` blob: the webhook reconciler writes @@ -4238,7 +4239,7 @@ async fn edit_deploy_ui_config( ) -> Result { require_admin(is_admin, &username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let args_for_audit = format!("{:?}", new_config.deploy_ui_settings); audit_log( &mut *tx, @@ -4304,7 +4305,7 @@ async fn edit_default_scripts( ) -> Result { require_admin(is_admin, &username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; audit_log( &mut *tx, @@ -4355,7 +4356,7 @@ async fn get_default_scripts( Extension(db): Extension, Path(w_id): Path, ) -> JsonResult> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let default_scripts = sqlx::query_scalar!( "SELECT default_scripts FROM workspace_settings WHERE workspace_id = $1", &w_id @@ -4386,7 +4387,7 @@ async fn edit_default_app( require_admin(is_admin, &username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let args_for_audit = format!("{:?}", new_config.default_app_path); audit_log( @@ -4474,7 +4475,7 @@ async fn edit_error_handler( // Normalize to new format (handles both old CLI and new CLI requests) let ee = ee.into_normalized(); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; if let Some(fallback_to_instance_alerts) = ee.fallback_to_instance_alerts { if fallback_to_instance_alerts { @@ -4608,7 +4609,7 @@ async fn edit_success_handler( // Normalize to new format (handles both old CLI and new CLI requests) let es = es.into_normalized(); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query_as!( Group, @@ -4687,7 +4688,7 @@ async fn set_environment_variable( ) -> Result { require_admin(authed.is_admin, &authed.username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; match value { Some(value) => { @@ -4821,7 +4822,7 @@ async fn set_encryption_key( // Build the previous cipher before the transaction (reads from cache/pool) let previous_encryption_key = build_crypt(&db, w_id.as_str()).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE workspace_key SET key = $1 WHERE workspace_id = $2", @@ -5030,7 +5031,7 @@ async fn user_workspaces( Extension(db): Extension, ApiAuthed { email, .. }: ApiAuthed, ) -> JsonResult { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let workspaces = sqlx::query_as!( UserWorkspace, "SELECT workspace.id, workspace.name, usr.username, workspace_settings.color, workspace.parent_workspace_id, @@ -5304,7 +5305,7 @@ async fn create_workspace( validate_workspace_name(&nw.name)?; - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; check_w_id_conflict(&mut tx, &nw.id).await?; sqlx::query!( @@ -7333,7 +7334,7 @@ async fn create_workspace_fork( ensure_no_existing_dev_workspace(&db, &parent_workspace_id).await?; } - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; if nw.is_dev_workspace { // The checks above ran outside a transaction, so the parent's eligibility and the chain's @@ -7660,7 +7661,7 @@ async fn attach_dev_workspace( ))); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Everything above ran outside a transaction, so prod's eligibility and the chain's labels could // have changed under us: re-decide both here, under the pairing lock. lock_dev_pairing(&mut tx, &[&prod_w_id, &dev_w_id]).await?; @@ -7873,7 +7874,7 @@ async fn detach_dev_workspace( let dev_w_id = req.dev_workspace_id; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Under the pairing lock, so a dev workspace cannot appear beneath this one between the check // below and the update. lock_dev_pairing(&mut tx, &[&prod_w_id, &dev_w_id]).await?; @@ -7965,7 +7966,7 @@ async fn edit_workspace( Json(ew): Json, ) -> Result { require_admin(is_admin, &username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE workspace SET name = $1, owner = $2 WHERE id = $3", ew.name, @@ -8002,7 +8003,7 @@ pub(crate) async fn archive_workspace_impl( dev_lock_parent: Option<&str>, ) -> Result<(usize, usize, usize)> { // Step 1: Disable all schedules and clear their queued jobs - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Unconditionally, before reading any pairing state: whether this workspace is a dev, and whether // it has one, is exactly what a concurrent attach changes, so gating the lock on the caller's // `dev_lock_parent` would skip it on the strength of the value the race invalidates. @@ -8172,7 +8173,7 @@ async fn archive_workspace( archive_workspace_impl(&db, &w_id, &authed.username, dev_lock_parent.as_deref()).await?; // Audit log - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let mut audit_params = HashMap::new(); audit_params.insert("disabled_schedules", schedules_count.to_string()); audit_params.insert("canceled_jobs", canceled_count.to_string()); @@ -8233,7 +8234,7 @@ async fn leave_workspace( authed: ApiAuthed, ) -> Result { windmill_api_auth::forbid_job_token_account_destruction(&authed)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "DELETE FROM usr WHERE workspace_id = $1 AND email = $2", &w_id, @@ -8272,7 +8273,7 @@ async fn unarchive_workspace( #[cfg(not(feature = "enterprise"))] _check_nb_of_workspaces(&db).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!("UPDATE workspace SET deleted = false WHERE id = $1", &w_id) .execute(&mut *tx) .await?; @@ -8343,7 +8344,7 @@ async fn invite_user( return Err(Error::BadRequest(msg)); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let already_in_workspace = sqlx::query_scalar!( "SELECT EXISTS (SELECT 1 FROM usr WHERE workspace_id = $1 AND email = $2)", @@ -8495,7 +8496,7 @@ async fn add_user( return Err(Error::BadRequest(msg)); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let already_exists_email = sqlx::query_scalar!( "SELECT EXISTS(SELECT 1 FROM usr WHERE workspace_id = $1 AND email = $2)", @@ -8659,7 +8660,7 @@ async fn delete_invite( ) -> Result<(StatusCode, String)> { require_admin(is_admin, &username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "DELETE FROM workspace_invite WHERE @@ -8862,7 +8863,7 @@ async fn change_workspace_name( ) -> Result { require_admin(authed.is_admin, &authed.username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE workspace SET name = $1 WHERE id = $2", @@ -8909,7 +8910,7 @@ async fn change_workspace_color( ) -> Result { require_admin(authed.is_admin, &authed.username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE workspace_settings SET color = $1 WHERE workspace_id = $2", @@ -9063,7 +9064,7 @@ async fn update_operator_settings( ) -> Result { require_admin(authed.is_admin, &authed.username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let settings_json = serde_json::json!(settings); @@ -9526,7 +9527,7 @@ async fn create_protection_rule( require_admin(authed.is_admin, &authed.username)?; reject_reserved_rule_name(&req.name)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Check if rule with this name already exists let exists = sqlx::query_scalar!( @@ -9622,7 +9623,7 @@ async fn update_protection_rule( reject_reserved_rule_name(new_name)?; } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Check if rule exists let exists = sqlx::query_scalar!( @@ -9721,7 +9722,7 @@ async fn delete_protection_rule( require_admin(authed.is_admin, &authed.username)?; reject_reserved_rule_name(&rule_name)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Delete the rule let result = sqlx::query!( @@ -10611,7 +10612,7 @@ async fn seed_full_diff_scan( let kinds: Vec = candidates.iter().map(|(k, _)| k.clone()).collect(); let paths: Vec = candidates.iter().map(|(_, p)| p.clone()).collect(); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; lock_workspace_pair(&mut tx, &w_id, &target_workspace_id).await?; // Re-checked under the lock: the pair may have been linked since the check above, // and `attach_dev_workspace` clears the arbitrary rows as it links. Seeding after diff --git a/backend/windmill-api-workspaces/src/workspaces_extra.rs b/backend/windmill-api-workspaces/src/workspaces_extra.rs index 1968437e40..7f390d3b45 100644 --- a/backend/windmill-api-workspaces/src/workspaces_extra.rs +++ b/backend/windmill-api-workspaces/src/workspaces_extra.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use std::collections::HashMap; use windmill_api_auth::{require_super_admin, ApiAuthed}; @@ -54,7 +55,7 @@ pub(crate) async fn change_workspace_id( require_admin(authed.is_admin, &authed.username)?; } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // A rename rewrites the workspace's dev flag and reparents its children, so it decides on the // same state the pairing handlers do: without this lock a concurrent create/attach could commit @@ -926,7 +927,7 @@ pub(crate) async fn delete_workspace( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; if !(is_fork && is_workspace_owner(&authed, &w_id, &mut tx).await?) && !windmill_api_auth::is_super_admin_authed(&db, &authed).await? { @@ -1294,7 +1295,7 @@ pub async fn drop_forked_datatable_databases( ) -> Result>> { // Same permission check as delete_workspace: fork owner or super admin let is_fork = workspace_is_fork(&db, &w_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; if !(is_fork && is_workspace_owner(&authed, &w_id, &mut tx).await?) && !windmill_api_auth::is_super_admin_authed(&db, &authed).await? { @@ -1451,7 +1452,7 @@ pub async fn drop_forked_ducklake_namespaces( Path(w_id): Path, ) -> Result>> { let is_fork = workspace_is_fork(&db, &w_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; if !(is_fork && is_workspace_owner(&authed, &w_id, &mut tx).await?) && !windmill_api_auth::is_super_admin_authed(&db, &authed).await? { diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 1a5f9dba87..ea278e47e8 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use crate::db::{ApiAuthed, DB}; use crate::utils::check_scopes; @@ -791,7 +792,7 @@ fn proxy_request_to_request_builder( } async fn audit_global_ai_request(db: &DB, authed: &ApiAuthed) -> Result<()> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; audit_log( &mut *tx, @@ -1169,7 +1170,7 @@ async fn proxy( // Handle GoogleAI (Gemini) using the native Gemini API if matches!(proxy_mode, ProxyExecutionMode::NativeGoogleAi) { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; audit_log( &mut *tx, &authed, @@ -1207,7 +1208,7 @@ async fn proxy( // Handle Bedrock-specific logic when the feature is enabled #[cfg(feature = "bedrock")] if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; audit_log( &mut *tx, &authed, @@ -1270,7 +1271,7 @@ async fn proxy( let response = request.send().await.map_err(to_anyhow)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; audit_log( &mut *tx, diff --git a/backend/windmill-api/src/ai_evals/results.rs b/backend/windmill-api/src/ai_evals/results.rs index 500af88d72..e9da1b4293 100644 --- a/backend/windmill-api/src/ai_evals/results.rs +++ b/backend/windmill-api/src/ai_evals/results.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use super::*; /// One run of a dataset: written once when the dataset is run, and only ever read afterwards. @@ -562,7 +563,7 @@ async fn resolve_deployed_draft( // Both writes in one transaction: a failure between them would leave the experiment promoted // to a version while its cells stayed a draft's, a split no later read repairs since the // experiment is no longer a draft. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE eval_experiment SET subject = jsonb_set( diff --git a/backend/windmill-api/src/ai_evals/run.rs b/backend/windmill-api/src/ai_evals/run.rs index e102abaf4c..d87f7f0c1a 100644 --- a/backend/windmill-api/src/ai_evals/run.rs +++ b/backend/windmill-api/src/ai_evals/run.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use super::*; /// Node id of the agent step. The answer is read back by this id, so it is part of the stored @@ -639,7 +640,7 @@ pub async fn run_experiment( // Recorded before the job is queued, so a launch that dies partway leaves an experiment naming // a job that never started rather than a flow no experiment accounts for and nothing collects. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let experiment_id = new_run( &mut tx, &w_id, diff --git a/backend/windmill-api/src/ai_skills.rs b/backend/windmill-api/src/ai_skills.rs index 11d99d55da..720ea36a01 100644 --- a/backend/windmill-api/src/ai_skills.rs +++ b/backend/windmill-api/src/ai_skills.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use crate::db::{ApiAuthed, DB}; use axum::{ extract::{Extension, Json, Path}, @@ -208,7 +209,7 @@ async fn upload_skills( } let names = collect_upload_names(&payload.skills)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let counts = sqlx::query!( r#"SELECT COUNT(*)::bigint AS "total!", @@ -268,7 +269,7 @@ async fn delete_skill( ) -> Result { require_admin(authed.is_admin, &authed.username)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let deleted = sqlx::query_scalar!( "DELETE FROM ai_skill WHERE workspace_id = $1 AND name = $2 RETURNING name", &w_id, diff --git a/backend/windmill-api/src/apps.rs b/backend/windmill-api/src/apps.rs index 0e64d1cba2..73cd08ed8c 100644 --- a/backend/windmill-api/src/apps.rs +++ b/backend/windmill-api/src/apps.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use std::{borrow::Cow, collections::HashMap, sync::Arc}; /* @@ -1366,7 +1367,7 @@ async fn mint_raw_app_sdk_token( // session lacks. Some(authed.read_only), ); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let token = create_token_internal(&mut *tx, db, authed, token_config).await?; tx.commit().await?; Ok((token, expiration)) @@ -1580,7 +1581,7 @@ pub async fn mint_app_embed_token( // session lacks. Some(authed.read_only), ); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let token = create_token_internal(&mut *tx, db, authed, token_config).await?; tx.commit().await?; Some((token, expiration)) diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index 513341b6c1..61b3e4f6f5 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -11,6 +11,7 @@ pub use windmill_api_jobs::query::*; pub use windmill_api_jobs::types::*; pub use windmill_api_sse::*; +use windmill_common::db::BeginCancelSafe; use axum::body::Body; use futures::{StreamExt, TryFutureExt}; use itertools::Itertools; @@ -551,7 +552,7 @@ async fn get_job_public_view_token( // The link is stateless and permanent, so the mint is the only moment this is // observable: audit it unconditionally rather than through the opt-in job-view log. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; audit_log( &mut *tx, &AuditAuthor::from(&authed), @@ -608,7 +609,7 @@ async fn cancel_job_api( require_job_update_read_access(&db, &user_db, authed, &w_id, &id, None).await?; } - let tx = db.begin().await?; + let tx = db.begin_cancel_safe().await?; let audit_author: AuditAuthor = match opt_authed.as_ref() { Some(authed) => (authed).into(), @@ -770,7 +771,7 @@ async fn force_cancel( require_job_update_read_access(&db, &user_db, authed, &w_id, &target, None).await?; } - let tx = db.begin().await?; + let tx = db.begin_cancel_safe().await?; let audit_author: AuditAuthor = match opt_authed.as_ref() { Some(authed) => (authed).into(), @@ -4507,7 +4508,7 @@ pub async fn resume_suspended_flow_as_owner( Path((w_id, flow_id)): Path<(String, Uuid)>, QueryOrBody(value): QueryOrBody, ) -> error::Result { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let (flow, job_id, is_wac) = get_suspended_flow_info(flow_id, &mut tx).await?; @@ -4624,7 +4625,7 @@ async fn resume_suspended( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Resolve the suspended flow (works for both WAC and classic flows) let (flow, resume_job_id, is_wac) = get_suspended_flow_info(job_id, &mut tx).await?; @@ -5132,7 +5133,7 @@ async fn resume_suspended_job_internal( } else { authed.as_ref().map(|x| x.username.clone()) }; - let mut tx: Transaction<'_, Postgres> = db.begin().await?; + let mut tx: Transaction<'_, Postgres> = db.begin_cancel_safe().await?; // Inside the transaction that inserts the row and moves the suspend counter: // validating earlier would let the workflow resolve this step and suspend on the @@ -5852,7 +5853,7 @@ pub async fn get_wac_approval_urls( // behind it and sees this key rather than racing past an earlier read. Upsert // because a workflow can mint before any step has checkpointed, and a bare // UPDATE would silently match nothing and leave the link unbound. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query( "INSERT INTO v2_job_status (id, workflow_as_code_status) VALUES ($1, jsonb_build_object('_minted_approval_keys', @@ -7281,7 +7282,7 @@ pub async fn wac_inline_checkpoint( })?; let source_hash = runnable_id.map(|h| h.to_string()); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let failure = windmill_common::wac::persist_inline_checkpoint_delta( &mut tx, &job_id, diff --git a/backend/windmill-api/src/offboarding.rs b/backend/windmill-api/src/offboarding.rs index a69441ea75..0b73ef1f0c 100644 --- a/backend/windmill-api/src/offboarding.rs +++ b/backend/windmill-api/src/offboarding.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use std::collections::HashMap; use crate::db::{ApiAuthed, OptJobAuthed}; @@ -427,7 +428,7 @@ pub(crate) async fn offboard_workspace_user( )) })?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let summary = offboard_user_from_workspace( &mut tx, @@ -564,7 +565,7 @@ pub(crate) async fn offboard_global_user( })); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let mut total_summary = OffboardSummary { scripts_reassigned: 0, diff --git a/backend/windmill-api/src/users.rs b/backend/windmill-api/src/users.rs index 8804004f2b..76b7aa8476 100644 --- a/backend/windmill-api/src/users.rs +++ b/backend/windmill-api/src/users.rs @@ -9,6 +9,7 @@ // Re-export everything from windmill-api-users pub use windmill_api_users::users::*; +use windmill_common::db::BeginCancelSafe; use std::sync::Arc; use crate::db::{ApiAuthed, OptJobAuthed}; @@ -184,7 +185,7 @@ async fn rename_user( require_super_admin(&db, &authed).await?; forbid_superadmin_job_token(&db, &authed.email, job_id).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let username_conflict = sqlx::query_scalar!( "SELECT EXISTS(SELECT 1 FROM usr WHERE username = $1 and email != $2 UNION SELECT 1 FROM password WHERE username = $1 UNION SELECT 1 FROM pending_user WHERE username = $1)", @@ -885,7 +886,7 @@ async fn reset_password( )); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Find the token and verify it's not expired let magic_link = sqlx::query!( diff --git a/backend/windmill-api/src/workspace_shared_ui.rs b/backend/windmill-api/src/workspace_shared_ui.rs index 810ea3420f..e1d429a34d 100644 --- a/backend/windmill-api/src/workspace_shared_ui.rs +++ b/backend/windmill-api/src/workspace_shared_ui.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use crate::db::{ApiAuthed, DB}; use axum::{ extract::{Extension, Json, Path}, @@ -156,7 +157,7 @@ async fn update_shared_ui( let files_json = serde_json::to_value(&payload.files) .map_err(|e| Error::internal_err(format!("serializing files: {e}")))?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( r#"INSERT INTO workspace_shared_ui (workspace_id, files, version, edited_at, edited_by) VALUES ($1, $2, 1, now(), $3) diff --git a/backend/windmill-api/src/workspaces.rs b/backend/windmill-api/src/workspaces.rs index 01eb32b7f6..5aeb4b4515 100644 --- a/backend/windmill-api/src/workspaces.rs +++ b/backend/windmill-api/src/workspaces.rs @@ -8,6 +8,7 @@ // Re-export everything from windmill-api-workspaces pub use windmill_api_workspaces::workspaces::*; +use windmill_common::db::BeginCancelSafe; use windmill_api_workspaces::workspaces::{build_copilot_settings_state, InstanceAISummary}; use crate::ai::{invalidate_ai_request_cache_for_workspace, AIConfig}; @@ -110,7 +111,7 @@ async fn edit_copilot_config( ai_config.validate_model_pricing()?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; sqlx::query!( "UPDATE workspace_settings SET ai_config = $1 WHERE workspace_id = $2", diff --git a/backend/windmill-common/src/bench.rs b/backend/windmill-common/src/bench.rs index 68c3cb5a31..668ccca66e 100644 --- a/backend/windmill-common/src/bench.rs +++ b/backend/windmill-common/src/bench.rs @@ -1,3 +1,4 @@ +use crate::db::BeginCancelSafe; use crate::{ worker::{write_file, WINDMILL_DIR}, DB, @@ -265,7 +266,7 @@ pub async fn benchmark_init(benchmark_jobs: i32, db: &DB) { .await .unwrap_or_else(|e| panic!("failed to clean up v2_job: {e:#}")); - let mut tx = db.begin().await.unwrap(); + let mut tx = db.begin_cancel_safe().await.unwrap(); match benchmark_kind.as_str() { "dedicated" => { // you need to create the script first, check https://github.com/windmill-labs/windmill/blob/b76a92cfe454c686f005c65f534e29e039f3c706/benchmarks/lib.ts#L47 diff --git a/backend/windmill-common/src/db.rs b/backend/windmill-common/src/db.rs index 0048abf438..0b2febad85 100644 --- a/backend/windmill-common/src/db.rs +++ b/backend/windmill-common/src/db.rs @@ -4,6 +4,42 @@ use crate::audit::AuditAuthor; pub type DB = Pool; +/// `Pool::begin` is not cancel-safe, and on Postgres it is the one place where losing the +/// race corrupts the pool rather than just the caller. +/// +/// sqlx sends `BEGIN` and only increments the transaction depth its rollback-on-drop guard +/// keys on once the round trip returns (`sqlx-postgres/src/transaction.rs`). A future +/// cancelled in between — a disconnecting API client, a `tokio::time::timeout`, an aborted +/// task — therefore leaves the session inside a transaction that nothing will ever roll +/// back, and sqlx hands that connection to the next borrower, whose statements then run +/// inside it and hold its locks. +/// +/// Upstream has known about it since 2022 (launchbadge/sqlx#2054) and fixed it for SQLite +/// only, so it is ours to work around. +pub trait BeginCancelSafe { + /// Begins a transaction that a cancelled caller cannot leak. + fn begin_cancel_safe( + &self, + ) -> impl std::future::Future, sqlx::Error>> + Send; +} + +impl BeginCancelSafe for Pool { + fn begin_cancel_safe( + &self, + ) -> impl std::future::Future, sqlx::Error>> + Send + { + let db = self.clone(); + async move { + // Dropping a `JoinHandle` detaches its task rather than aborting it, so a + // cancelled caller still leaves the `Transaction` to be built and then dropped + // — and dropping one at depth 1 is what queues the `ROLLBACK`. + tokio::spawn(async move { db.begin().await }) + .await + .map_err(|_| sqlx::Error::WorkerCrashed)? + } + } +} + /// Workspace ID resolved by gateway middleware (stored in request extensions). /// Used by auth to resolve workspace when the URL path doesn't contain one. #[derive(Clone, Debug)] @@ -187,7 +223,7 @@ impl<'c, 'd, T: Authable + Sync> Acquire<'c> for &'c DbWithOptAuthed<'d, T> { DbWithOptAuthed::UserDB { authed, user_db, .. } => { user_db.clone().begin(&**authed).await } - DbWithOptAuthed::DB { db, .. } => db.clone().begin().await, + DbWithOptAuthed::DB { db, .. } => db.clone().begin_cancel_safe().await, } }) } @@ -200,7 +236,7 @@ impl<'c, 'd, T: Authable + Sync> Acquire<'c> for &'c DbWithOptAuthed<'d, T> { DbWithOptAuthed::UserDB { authed, user_db, .. } => { user_db.clone().begin(&**authed).await } - DbWithOptAuthed::DB { db, .. } => db.clone().begin().await, + DbWithOptAuthed::DB { db, .. } => db.clone().begin_cancel_safe().await, } }) } @@ -230,7 +266,7 @@ impl UserDB { // folders_write // ); - let mut tx = self.db.begin().await?; + let mut tx = self.db.begin_cancel_safe().await?; if let Some(schema) = PG_SCHEMA.as_ref() { // SAFETY: `schema` is an operator-controlled environment variable (PG_SCHEMA), set at deploy time and never user-supplied. diff --git a/backend/windmill-common/src/dbt_manifest.rs b/backend/windmill-common/src/dbt_manifest.rs index 8fb124a33f..306f9e7e48 100644 --- a/backend/windmill-common/src/dbt_manifest.rs +++ b/backend/windmill-common/src/dbt_manifest.rs @@ -53,6 +53,7 @@ //! for the script it just ran. A new call site that cannot name where it //! authorized is a bug. +use crate::db::BeginCancelSafe; use std::collections::{BTreeMap, HashMap}; use serde::{Deserialize, Serialize}; @@ -300,7 +301,7 @@ pub async fn prune_dbt_run_graphs( // In ONE transaction with the orphan sweep: a restart in the gap leaves graph // rows whose marker is gone, and since the sweep runs only when a marker went, // every later call computes `retired == 0` and skips them for good. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Ordered by the script's own `created_at`, so "newest" is the newest deploy // and not the newest ingest: a late-finishing job must not promote an old // version. Scoped to one (workspace, path), which the path index serves. diff --git a/backend/windmill-common/src/scripts.rs b/backend/windmill-common/src/scripts.rs index 31219621e9..44855c678f 100644 --- a/backend/windmill-common/src/scripts.rs +++ b/backend/windmill-common/src/scripts.rs @@ -8,6 +8,7 @@ pub use windmill_types::scripts::*; +use crate::db::BeginCancelSafe; use crate::{ error::{to_anyhow, Error}, runnable_settings::{self}, @@ -369,7 +370,7 @@ pub async fn clone_script<'c>( deployment_message: Option, db: &DB, ) -> crate::error::Result { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let s = if let Some(s) = fetch_script_for_update(path, w_id, &mut *tx).await? { s } else { diff --git a/backend/windmill-common/src/usernames.rs b/backend/windmill-common/src/usernames.rs index 6fde2a8186..ed671e3020 100644 --- a/backend/windmill-common/src/usernames.rs +++ b/backend/windmill-common/src/usernames.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use crate::db::BeginCancelSafe; use regex::Regex; use sqlx::{Postgres, Transaction}; @@ -59,7 +60,7 @@ pub async fn generate_instance_wide_unique_username<'c>( } pub async fn generate_instance_username_for_all_users(db: &DB) -> error::Result<()> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let users = sqlx::query!(r#"SELECT p.email as "email!", u.username as "username?" FROM password p LEFT JOIN usr u ON p.email = u.email WHERE p.username IS NULL AND (SELECT COUNT(DISTINCT username) FROM usr WHERE email = p.email) <= 1"#) .fetch_all(&mut *tx) .await?; diff --git a/backend/windmill-common/src/utils.rs b/backend/windmill-common/src/utils.rs index 62d47a34d7..3294831853 100644 --- a/backend/windmill-common/src/utils.rs +++ b/backend/windmill-common/src/utils.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use crate::db::BeginCancelSafe; use crate::auth::is_devops_email; use crate::ee_oss::LICENSE_KEY_ID; #[cfg(feature = "enterprise")] @@ -1426,7 +1427,7 @@ pub async fn get_custom_pg_instance_replication_password(db: &DB) -> Result) { + // One connection, so the session inspected below is the one the cancelled begin used. + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .min_connections(0) + .connect_with((*db.connect_options()).clone()) + .await + .expect("failed to build pool"); + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&pool) + .await + .unwrap(); + + // Polls exactly once — enough to start the work, never enough to finish it — and then + // drops the future, so the cancellation lands while the transaction is being opened. A + // short timeout would not do: tokio's timer granularity lets the begin win the race. + let cancelled = pool.begin_cancel_safe().now_or_never(); + assert!(cancelled.is_none(), "the begin must not have completed"); + + // The detached task still finishes and drops its transaction, which rolls it back. + tokio::time::sleep(Duration::from_millis(500)).await; + + let mut admin = PgConnection::connect_with(&(*db.connect_options()).clone()) + .await + .expect("failed to open an observing connection"); + let (state, last_query): (Option, Option) = + sqlx::query_as("SELECT state, query FROM pg_stat_activity WHERE pid = $1") + .bind(pid) + .fetch_one(&mut admin) + .await + .unwrap(); + assert_eq!( + state.as_deref(), + Some("idle"), + "session left in a transaction" + ); + // Proves the task outlived its caller rather than being aborted with it: it opened the + // transaction and rolled it back. Cancelling `Pool::begin` directly leaves the session + // still showing whatever ran before it. + assert_eq!( + last_query.as_deref(), + Some("ROLLBACK"), + "the cancelled begin should still have opened and rolled back a transaction" + ); + + sqlx::query_scalar::<_, i32>("SELECT 1") + .fetch_one(&pool) + .await + .expect("pool must still serve queries"); +} diff --git a/backend/windmill-dep-map/src/lib.rs b/backend/windmill-dep-map/src/lib.rs index ee6a3c3fd4..0608565165 100644 --- a/backend/windmill-dep-map/src/lib.rs +++ b/backend/windmill-dep-map/src/lib.rs @@ -5,6 +5,7 @@ pub mod scoped_dependency_map; pub mod trigger_dependents; pub mod workspace_dependencies; +use windmill_common::db::BeginCancelSafe; use std::collections::HashMap; use std::path::{Component, Path, PathBuf}; @@ -146,7 +147,7 @@ pub async fn process_relative_imports( // TODO: Should be moved into handle_dependency_job body to be more consistent with how flows and apps are handled { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let mut dependency_map = ScopedDependencyMap::fetch_maybe_rearranged( &w_id, script_path, diff --git a/backend/windmill-dep-map/src/scoped_dependency_map.rs b/backend/windmill-dep-map/src/scoped_dependency_map.rs index ad1093bb76..ced3da6be7 100644 --- a/backend/windmill-dep-map/src/scoped_dependency_map.rs +++ b/backend/windmill-dep-map/src/scoped_dependency_map.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use serde::Serialize; use sqlx::PgExecutor; use tokio::sync::RwLock; @@ -422,7 +423,7 @@ SELECT importer_node_id, imported_path, imported_lockfile_hash { let (sd, smd) = cache::script::fetch(&db.clone().into(), r.hash.into()).await?; let mut dmap = ScopedDependencyMap::fetch(w_id, &r.path, "script", db).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; tx = dmap .patch( @@ -446,7 +447,7 @@ SELECT importer_node_id, imported_path, imported_lockfile_hash let mut dmap = ScopedDependencyMap::fetch(w_id, &r.path, "flow", db).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let mut to_process = vec![]; let flow_value = flow_data.value(); let mut modules_to_check = flow_value.modules.iter().collect::>(); @@ -504,7 +505,7 @@ SELECT importer_node_id, imported_path, imported_lockfile_hash .await?; let mut dmap = ScopedDependencyMap::fetch(w_id, &r.path, "app", db).await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let mut to_process = vec![]; traverse_app_inline_scripts(&value, None, &mut |ais, id| { to_process.push(( diff --git a/backend/windmill-dep-map/src/workspace_dependencies.rs b/backend/windmill-dep-map/src/workspace_dependencies.rs index 03db1405e5..5a4600243b 100644 --- a/backend/windmill-dep-map/src/workspace_dependencies.rs +++ b/backend/windmill-dep-map/src/workspace_dependencies.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use serde::{Deserialize, Serialize}; use windmill_common::{ cache::workspace_dependencies::EXISTS_CACHE_TIMEOUT, error, scripts::ScriptLang, @@ -71,7 +72,7 @@ impl NewWorkspaceDependencies { } }; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let prev_description = sqlx::query_scalar!( " UPDATE workspace_dependencies diff --git a/backend/windmill-native-triggers/src/google/external.rs b/backend/windmill-native-triggers/src/google/external.rs index 1c5a2dd9ec..a0988ed7f9 100644 --- a/backend/windmill-native-triggers/src/google/external.rs +++ b/backend/windmill-native-triggers/src/google/external.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use async_trait::async_trait; use reqwest::Method; use serde_json::value::RawValue; @@ -568,7 +569,7 @@ async fn try_renew_channel_locked( .renew_channel(workspace_id, &fresh_trigger, db) .await?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Past this point a new Google channel exists. Any failure leaks it. if let Err(e) = update_native_trigger_service_config( diff --git a/backend/windmill-native-triggers/src/rename.rs b/backend/windmill-native-triggers/src/rename.rs index a7168f5c7e..b3be8da653 100644 --- a/backend/windmill-native-triggers/src/rename.rs +++ b/backend/windmill-native-triggers/src/rename.rs @@ -1,5 +1,6 @@ //! Keeping native triggers usable when their runnable is renamed. +use windmill_common::db::BeginCancelSafe; use std::collections::HashMap; use windmill_api_auth::ApiAuthed; @@ -140,7 +141,7 @@ async fn reregister_one( // The token is scoped to the runnable path and only its hash is kept, so pointing the webhook // at the new path means minting a replacement rather than reusing the old one. Commit it // before handing it out: a service may call back the moment it accepts the new URL. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let webhook_token = new_webhook_token( &mut tx, db, @@ -177,7 +178,7 @@ async fn reregister_one( ) })?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let applied = record_reregistration( &mut *tx, diff --git a/backend/windmill-native-triggers/src/workspace_integrations.rs b/backend/windmill-native-triggers/src/workspace_integrations.rs index eccfa56f19..e1f8a43788 100644 --- a/backend/windmill-native-triggers/src/workspace_integrations.rs +++ b/backend/windmill-native-triggers/src/workspace_integrations.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use axum::Router; #[cfg(feature = "native_trigger")] @@ -430,7 +431,7 @@ async fn list_integrations( Path(workspace_id): Path, ) -> JsonResult> { require_admin(authed.is_admin, &workspace_id)?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let integrations = sqlx::query_as!( WorkspaceIntegrations, r#" diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index a1bd8034f0..87cb261d5e 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use std::future::Future; use std::time::Duration; use std::{collections::HashMap, sync::Arc, vec}; @@ -595,7 +596,7 @@ async fn cancel_persistent_script_jobs_internal<'c>( w_id: &str, db: &Pool, ) -> error::Result> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // we could have retrieved the job IDs in the first query where we retrieve the hashes, but just in case a job was inserted in the queue right in-between the two above query, we re-do the fetch here let jobs_to_cancel = sqlx::query_scalar::<_, Uuid>( @@ -1160,7 +1161,7 @@ async fn commit_completed_job( let serialized_result = result.serialized_json(); let sanitized_result = strip_json_nul(serialized_result.as_ref()); - let mut tx = db.begin().warn_after_seconds(10).await?; + let mut tx = db.begin_cancel_safe().warn_after_seconds(10).await?; let duration = sqlx::query_scalar!( "INSERT INTO v2_job_completed AS cj @@ -3824,7 +3825,7 @@ impl PulledJobResult { // otherwise let a zombie re-pull see its own prior claim and keep only // its own args (dropping the siblings it had claimed). One transaction // makes the claim and the merged-args write commit together (or neither). - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Emitted AFTER the transaction commits — writing logs via a second pool // connection while the claim tx + row locks are held risks pool-exhaustion // stalls under concurrent debounced pulls. @@ -5187,7 +5188,7 @@ impl<'c> PushIsolationLevel<'c> { pub async fn into_tx(self) -> error::Result> { match self { PushIsolationLevel::Isolated(db, authed) => Ok((db.begin(&authed).await?).into()), - PushIsolationLevel::IsolatedRoot(db) => Ok(db.begin().await?), + PushIsolationLevel::IsolatedRoot(db) => Ok(db.begin_cancel_safe().await?), PushIsolationLevel::Transaction(tx) => Ok(tx), } } diff --git a/backend/windmill-queue/src/schedule.rs b/backend/windmill-queue/src/schedule.rs index fa495c9cd2..1eb13b9720 100644 --- a/backend/windmill-queue/src/schedule.rs +++ b/backend/windmill-queue/src/schedule.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use crate::push; use crate::PushIsolationLevel; use anyhow::Context; @@ -636,7 +637,7 @@ pub enum RearmOutcome { /// callers (the monitor's reconciliation pass) only. A caller acting for a user /// MUST already have enforced their permissions on `w_id` and `path`. pub async fn rearm_schedule(db: &DB, w_id: &str, path: &str) -> Result { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Lock the row for the whole push: an edit or a disable committing between the // read and the push would otherwise leave a queued occurrence for a schedule // that is disabled, or one built from superseded settings. diff --git a/backend/windmill-store/src/resources.rs b/backend/windmill-store/src/resources.rs index b107d34be4..b43d4395b3 100644 --- a/backend/windmill-store/src/resources.rs +++ b/backend/windmill-store/src/resources.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use dashmap::DashMap; use std::collections::HashMap; use std::net::IpAddr; @@ -2419,7 +2420,7 @@ async fn clear_resource_history( return Err(Error::NotFound(format!("Resource {} not found", path))); } - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let deleted = sqlx::query!( "DELETE FROM resource_version rv WHERE rv.workspace_id = $1 AND rv.path = $2 diff --git a/backend/windmill-trigger/src/listener.rs b/backend/windmill-trigger/src/listener.rs index d2c80b4d65..a3ee87760a 100644 --- a/backend/windmill-trigger/src/listener.rs +++ b/backend/windmill-trigger/src/listener.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use std::{collections::HashMap, fmt::Debug, sync::Arc}; use crate::{ @@ -390,7 +391,7 @@ pub trait Listener: TriggerCrud + TriggerJobArgs { // lock spans both writes. let mut history_err = None; let report_status = async { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // SAFETY: Self::TABLE_NAME is a compile-time constant. let rows = sqlx::query(&format!( r#" diff --git a/backend/windmill-worker/src/ai/tools.rs b/backend/windmill-worker/src/ai/tools.rs index 9c15229ba6..dfb1ecd8e6 100644 --- a/backend/windmill-worker/src/ai/tools.rs +++ b/backend/windmill-worker/src/ai/tools.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use crate::ai::stream_event_processor::StreamEventProcessor; use crate::ai::utils::{ add_message_to_conversation, execute_mcp_tool, get_step_name_from_flow, @@ -459,7 +460,7 @@ async fn execute_windmill_tool( } }; - let mut tx = ctx.db.begin().await?; + let mut tx = ctx.db.begin_cancel_safe().await?; let job_perms = windmill_common::auth::get_job_perms(&mut *tx, &ctx.job.id, &ctx.job.workspace_id) diff --git a/backend/windmill-worker/src/ai/utils.rs b/backend/windmill-worker/src/ai/utils.rs index 6e80551650..b6ebaefc25 100644 --- a/backend/windmill-worker/src/ai/utils.rs +++ b/backend/windmill-worker/src/ai/utils.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use anyhow::Context; use serde_json::value::RawValue; use sqlx::types::Json; @@ -210,7 +211,7 @@ pub async fn add_message_to_conversation( step_name: &Option, success: bool, ) -> Result<(), Error> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; add_message_to_conversation_tx( &mut tx, *conversation_id, diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index d948fee088..b0e3602667 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use crate::ai::tools::{execute_tool_calls, ToolAbortHandles, ToolExecutionContext}; use crate::ai::utils::{ add_message_to_conversation, any_tool_needs_previous_result, cleanup_mcp_clients, @@ -2034,7 +2035,7 @@ async fn cleanup_orphaned_tool_jobs( } }; - let tx = match db.begin().await { + let tx = match db.begin_cancel_safe().await { Ok(tx) => tx, Err(e) => { tracing::error!( diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index 1fb566e57f..588788d5e2 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; #[cfg(feature = "deno_core")] use std::time::Instant; use std::{ @@ -2813,7 +2814,7 @@ pub async fn handle_wac_v2_output( // Step 1: Save checkpoint, suspend parent, and seed child checkpoints // in a single transaction — all BEFORE children become visible. { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Update checkpoint with pending steps update_checkpoint_for_dispatch(&mut checkpoint, &steps, &mode, &job_ids); @@ -3202,7 +3203,7 @@ pub async fn handle_wac_v2_output( job_ids: serde_json::Map::new(), }); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Save checkpoint let status_json = serde_json::to_value(&checkpoint).map_err(|e| { @@ -3397,7 +3398,7 @@ pub async fn handle_wac_v2_output( job_ids: serde_json::Map::new(), }); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Save checkpoint let status_json = serde_json::to_value(&checkpoint).map_err(|e| { @@ -3492,7 +3493,7 @@ pub async fn handle_wac_v2_output( // into `persist_inline_checkpoint_delta` preserves the original // atomicity from before the shared-helper refactor. let source_hash = job.runnable_id.map(|h| h.0.to_string()); - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; crate::wac_executor::persist_inline_checkpoint_delta( &mut tx, diff --git a/backend/windmill-worker/src/dbt_executor.rs b/backend/windmill-worker/src/dbt_executor.rs index 33ddd6c32e..5595e36e0c 100644 --- a/backend/windmill-worker/src/dbt_executor.rs +++ b/backend/windmill-worker/src/dbt_executor.rs @@ -8,6 +8,7 @@ //! while the run is in flight, and the structured job result comes from //! `run_results.json` at the end. +use windmill_common::db::BeginCancelSafe; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::process::Stdio; @@ -688,7 +689,7 @@ pub(crate) async fn dbt_dep( // No warehouse identity, so nothing can be ingested — but this version's // rows must still go, or a descriptor moved to its own profiles.yml keeps // claiming relations it no longer describes. - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Clearing is a publication too: an older job that no longer describes the // script must not wipe a newer deploy's graph. let published = claim_graph_publication(&mut tx, w_id, script_path, publisher).await?; @@ -3104,7 +3105,7 @@ async fn run_parse_only( result.graph_job = stored.then_some(job.id); } None => { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; windmill_common::dbt_manifest::replace_dbt_editor_graph( &mut tx, &job.workspace_id, @@ -3268,7 +3269,7 @@ async fn persist_ingest( // deploys nothing and must not touch what a deploy wrote. return Ok(false); }; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Deletion only soft-updates `script`, so re-inserting would republish model // SQL a user deleted. NOT `archived`: every redeploy archives the parent, and // treating that as deletion would deny v1 the graph its own runs render. An diff --git a/backend/windmill-worker/src/duckdb_executor.rs b/backend/windmill-worker/src/duckdb_executor.rs index 54457f81b0..d58cdafa28 100644 --- a/backend/windmill-worker/src/duckdb_executor.rs +++ b/backend/windmill-worker/src/duckdb_executor.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use std::cell::RefCell; use std::env; use std::ffi::{c_char, c_uint, CStr, CString}; @@ -1280,7 +1281,7 @@ async fn record_asset_schema_best_effort( snapshot_id: Option, job_id: Uuid, ) -> anyhow::Result<()> { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; windmill_common::materialization::record_asset_schema( &mut tx, w_id, diff --git a/backend/windmill-worker/src/result_processor.rs b/backend/windmill-worker/src/result_processor.rs index 647d569cbe..8cc5ee6bd6 100644 --- a/backend/windmill-worker/src/result_processor.rs +++ b/backend/windmill-worker/src/result_processor.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; #[cfg(feature = "otel")] use opentelemetry::trace::FutureExt; @@ -1966,7 +1967,7 @@ pub(crate) async fn handle_wac_child_completion( let result_json = serde_json::to_value(&result_value) .map_err(|e| error::Error::InternalErr(format!("Failed to serialize step result: {e}")))?; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Merge the completed step into the checkpoint. // Uses `|| jsonb_build_object(key, value)` so concurrent children on diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index bd39afa8c7..de4fc61dce 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use windmill_common::db::BeginCancelSafe; use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; @@ -756,7 +757,7 @@ pub async fn update_flow_status_after_job_completion_internal( _ => false, }; - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; add_time!(bench, "process module status START"); @@ -2245,7 +2246,7 @@ async fn add_tool_message_to_conversation( }; // Insert new assistant message - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; add_message_to_conversation_tx( &mut tx, conversation_id, @@ -2884,7 +2885,7 @@ pub async fn handle_flow( let mini_job = MiniCompletedJob::from(flow_job.clone()); let runnable_path = flow_job.runnable_path.as_ref().unwrap().clone(); let schedule_push_result = (|| async { - let tx = db.begin().warn_after_seconds(5).await + let tx = db.begin_cancel_safe().warn_after_seconds(5).await .map_err(|e| Error::internal_err(format!("begin tx for schedule push: {e:#}")))?; let (tx, schedule_push_err) = try_schedule_next_job( db, @@ -2929,7 +2930,7 @@ pub async fn handle_flow( // so once the flow is gone nothing reaches this code again. let mut history_lost = None; let disable_result = async { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let rows = sqlx::query!( "UPDATE schedule SET enabled = false, error = $1 WHERE workspace_id = $2 AND path = $3 AND enabled = true", err.to_string(), @@ -3407,7 +3408,7 @@ async fn push_next_flow_job( FlowStatusModule::WaitingForPriorSteps { .. } | FlowStatusModule::WaitingForEvents { .. } ) { if let Some((suspend, last)) = needs_resume(&flow, &status) { - let mut tx = db.begin().warn_after_seconds(3).await?; + let mut tx = db.begin_cancel_safe().warn_after_seconds(3).await?; /* Lock this row to prevent the suspend column getting out out of sync * if a resume message arrives after we fetch and count them here. @@ -4169,7 +4170,7 @@ async fn push_next_flow_job( }; let len = job_payloads.len(); - let mut tx = db.begin().warn_after_seconds(3).await?; + let mut tx = db.begin_cancel_safe().warn_after_seconds(3).await?; let nargs = args.as_ref(); for (i, payload_tag) in job_payloads.into_iter().enumerate() { if i % 100 == 0 && i != 0 { @@ -4915,7 +4916,7 @@ async fn push_next_flow_job( // same_worker_tx: Sender, // base_internal_url: &str, // ) -> error::Result<()> { -// let mut tx = db.begin().await?; +// let mut tx = db.begin_cancel_safe().await?; // let next_step = i // .checked_add(1) diff --git a/backend/windmill-worker/src/worker_lockfiles.rs b/backend/windmill-worker/src/worker_lockfiles.rs index 5cfb9f08ce..6e1a8163f6 100644 --- a/backend/windmill-worker/src/worker_lockfiles.rs +++ b/backend/windmill-worker/src/worker_lockfiles.rs @@ -1,3 +1,4 @@ +use windmill_common::db::BeginCancelSafe; use std::borrow::Cow; use std::collections::HashMap; use std::fs::{create_dir_all, remove_dir_all}; @@ -755,7 +756,7 @@ pub async fn handle_flow_dependency_job( let mut dependency_map = match tokio::time::timeout( std::time::Duration::from_secs(15), async { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; let mut dependency_map = ScopedDependencyMap::fetch_maybe_rearranged( &job.workspace_id, @@ -917,7 +918,7 @@ pub async fn handle_flow_dependency_job( Done, } let phase3 = tokio::time::timeout(std::time::Duration::from_secs(30), async { - let mut tx = db.begin().await?; + let mut tx = db.begin_cancel_safe().await?; // Counterpart to phase 1's conditional dissolve. if triggered_by_relative_import { @@ -2142,7 +2143,7 @@ async fn lock_modules_app( .patch( referenced_paths.clone(), container_id.unwrap_or_default(), - db.begin().await?, + db.begin_cancel_safe().await?, ) .await? .commit() @@ -2161,7 +2162,7 @@ async fn lock_modules_app( .patch( referenced_paths.clone(), container_id.unwrap_or_default(), - db.begin().await?, + db.begin_cancel_safe().await?, ) .await? .commit() @@ -2207,7 +2208,7 @@ async fn lock_modules_app( .patch( referenced_paths.clone(), container_id.unwrap_or_default(), - db.begin().await?, + db.begin_cancel_safe().await?, ) .await? .commit() @@ -2442,7 +2443,7 @@ pub async fn handle_app_dependency_job( // TODO: Dissolve in the end? dependency_map - .dissolve(db.begin().await?) + .dissolve(db.begin_cancel_safe().await?) .await .commit() .await?;