starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-09-21 00:02:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: make Pool::begin cancel-safe across the backend

Browse Source 
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Qzqmh52NU8fB9RBQNNkJGt
This commit is contained in:
Ruben Fiszel
2026-08-24 17:34:45 +00:00
co-authored by Claude Opus 5
parent 9c557859c5
commit dead8df080
51 changed files with 325 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/main.rs
+2 -1
View File
@@ -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)")
backend/src/monitor.rs
+8 -7
View File
@@ -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<Postgres>) {
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<DateTime<Utc>>)> {
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<Postgres>) {
// 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<Utc>,
) -> 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,
backend/windmill-api-configs/src/lib.rs
+3 -2
View File
@@ -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<String> {
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)
backend/windmill-api-flows/src/flows.rs
+2 -1
View File
@@ -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<DeploymentStatus> {
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
backend/windmill-api-groups/src/groups.rs
+10 -9
View File
@@ -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<IGroupUpdate>,
) -> Result<String> {
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<String>,
) -> Result<String> {
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<String> {
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<DB>) -> JsonResult<Vec<IGroup>> {
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<DB>) -> JsonResult<Vec<IGroup>> {
async fn list_igroups_with_workspaces(
Extension(db): Extension<DB>,
) -> JsonResult<Vec<IGroupWithWorkspaces>> {
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<Email>,
) -> Result<String> {
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<DB>,
) -> JsonResult<Vec<ExportedIGroup>> {
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<Vec<ExportedIGroup>>,
) -> Result<String> {
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
backend/windmill-api-jobs/src/concurrency_groups.rs
+2 -1
View File
@@ -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",
backend/windmill-api-jobs/src/execution.rs
+4 -3
View File
@@ -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<Vec<Uuid>> {
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()),
backend/windmill-api-schedule/src/lib.rs
+3 -2
View File
@@ -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<DB>,
Path((w_id, path)): Path<(String, StripPath)>,
) -> JsonResult<bool> {
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<String>;
match payload.handler_type {
HandlerType::Error => {
backend/windmill-api-scripts/src/scripts.rs
+5 -4
View File
@@ -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<DB>,
Path((w_id, hash_str)): Path<(String, String)>,
) -> Result<String> {
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<DB>,
Path((w_id, hash)): Path<(String, ScriptHash)>,
) -> JsonResult<DeploymentStatus> {
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 {
backend/windmill-api-settings/src/log_cleanup.rs
+2 -1
View File
@@ -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<String>, Option<DateTime<Utc>>)> {
let mut tx = db.begin().await?;
let mut tx = db.begin_cancel_safe().await?;
let active_root_job_ids: Vec<Uuid> = sqlx::query_scalar!(
"SELECT q.id FROM v2_job_queue q
backend/windmill-api-users/src/users.rs
+25 -24
View File
@@ -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<DB>,
) -> JsonResult<Vec<WorkspaceInvite>> {
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<NewInviteCode>,
// ) -> 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<DB>,
Json(nu): Json<DeclineInvite>,
) -> 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<DB>,
Json(nu): Json<AcceptInvite>,
) -> 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<DB>, authed: ApiAuthed) -> Result<String> {
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<EditWorkspaceUser>,
) -> Result<String> {
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<String> {
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<String> {
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<bool> = None;
if let Some(sa) = eu.is_super_admin {
@@ -1653,7 +1654,7 @@ async fn delete_user(
) -> Result<String> {
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<DB>,
Path((w_id, username_to_delete)): Path<(String, String)>,
) -> Result<String> {
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<String> {
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<String>,
) -> Result<String> {
forbid_job_token_account_destruction(&authed)?;
let mut tx = db.begin().await?;
let mut tx = db.begin_cancel_safe().await?;
let tokens_deleted: Vec<String> = 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<String> = 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<String> {
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<DB>,
) -> JsonResult<InstanceUsernameInfo> {
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<Vec<ExportedGlobalUser>> {
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<String> {
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?;
backend/windmill-api-workspaces/src/datatable_migrations.rs
+6 -5
View File
@@ -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<String> {
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();
backend/windmill-api-workspaces/src/deployment_requests.rs
+5 -4
View File
@@ -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
backend/windmill-api-workspaces/src/workspaces.rs
+40 -39
View File
@@ -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<EditCommandScript>,
) -> Result<String> {
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<String> {
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<String> {
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<String, DataTable> = 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<String> {
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<String> {
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<DB>,
Path(w_id): Path<String>,
) -> JsonResult<Option<serde_json::Value>> {
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<String> {
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<DB>,
ApiAuthed { email, .. }: ApiAuthed,
) -> JsonResult<WorkspaceList> {
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<EditWorkspace>,
) -> Result<String> {
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<String> {
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<String> {
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<String> {
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<String> {
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<String> = candidates.iter().map(|(k, _)| k.clone()).collect();
let paths: Vec<String> = 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
backend/windmill-api-workspaces/src/workspaces_extra.rs
+5 -4
View File
@@ -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<Json<Vec<String>>> {
// 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<String>,
) -> Result<Json<Vec<String>>> {
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?
{
backend/windmill-api/src/ai.rs
+5 -4
View File
@@ -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,
backend/windmill-api/src/ai_evals/results.rs
+2 -1
View File
@@ -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(
backend/windmill-api/src/ai_evals/run.rs
+2 -1
View File
@@ -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,
backend/windmill-api/src/ai_skills.rs
+3 -2
View File
@@ -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<String> {
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,
backend/windmill-api/src/apps.rs
+3 -2
View File
@@ -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))
backend/windmill-api/src/jobs.rs
+9 -8
View File
@@ -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<serde_json::Value>,
) -> error::Result<StatusCode> {
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,
backend/windmill-api/src/offboarding.rs
+3 -2
View File
@@ -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,
backend/windmill-api/src/users.rs
+3 -2
View File
@@ -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!(
backend/windmill-api/src/workspace_shared_ui.rs
+2 -1
View File
@@ -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)
backend/windmill-api/src/workspaces.rs
+2 -1
View File
@@ -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",
backend/windmill-common/src/bench.rs
+2 -1
View File
@@ -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
backend/windmill-common/src/db.rs
+39 -3
View File
@@ -4,6 +4,42 @@ use crate::audit::AuditAuthor;
pub type DB = Pool<Postgres>;
/// `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<Output = Result<Transaction<'static, Postgres>, sqlx::Error>> + Send;
}
impl BeginCancelSafe for Pool<Postgres> {
fn begin_cancel_safe(
&self,
) -> impl std::future::Future<Output = Result<Transaction<'static, Postgres>, 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.
backend/windmill-common/src/dbt_manifest.rs
+2 -1
View File
@@ -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.
backend/windmill-common/src/scripts.rs
+2 -1
View File
@@ -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<String>,
db: &DB,
) -> crate::error::Result<ClonedScript> {
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 {
backend/windmill-common/src/usernames.rs
+2 -1
View File
@@ -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?;
backend/windmill-common/src/utils.rs
+2 -1
View File
@@ -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<Stri
// serialize concurrent workers: otherwise two callers both rotate, and the second
// rotation invalidates the password the first already returned. Rotating and reading in
// one locked transaction keeps the decision atomic.
let mut tx = db.begin().await?;
let mut tx = db.begin_cancel_safe().await?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtext('custom_instance_replication_pwd'))")
.execute(&mut *tx)
.await?;
backend/windmill-common/tests/begin_cancel_safe.rs
+60
View File
@@ -0,0 +1,60 @@
//! Pins that a caller cancelled while `begin_cancel_safe` is in flight leaves no transaction
//! open on the session. sqlx's own `Pool::begin` does not hold this: its rollback-on-drop
//! guard keys on a transaction depth it only raises after the `BEGIN` round trip returns.
use futures::FutureExt;
use sqlx::{Connection, PgConnection, Pool, Postgres};
use std::time::Duration;
use windmill_common::db::BeginCancelSafe;
#[sqlx::test(migrations = "../migrations")]
async fn cancelled_begin_leaves_no_open_transaction(db: Pool<Postgres>) {
// 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<String>, Option<String>) =
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");
}
backend/windmill-dep-map/src/lib.rs
+2 -1
View File
@@ -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,
backend/windmill-dep-map/src/scoped_dependency_map.rs
+4 -3
View File
@@ -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::<Vec<_>>();
@@ -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((
backend/windmill-dep-map/src/workspace_dependencies.rs
+2 -1
View File
@@ -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
backend/windmill-native-triggers/src/google/external.rs
+2 -1
View File
@@ -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(
backend/windmill-native-triggers/src/rename.rs
+3 -2
View File
@@ -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<T: External>(
// 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<T: External>(
)
})?;
let mut tx = db.begin().await?;
let mut tx = db.begin_cancel_safe().await?;
let applied = record_reregistration(
&mut *tx,
backend/windmill-native-triggers/src/workspace_integrations.rs
+2 -1
View File
@@ -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<String>,
) -> JsonResult<Vec<WorkspaceIntegrations>> {
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#"
backend/windmill-queue/src/jobs.rs
+5 -4
View File
@@ -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<Postgres>,
) -> error::Result<Vec<Uuid>> {
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<T: Serialize + Send + Sync + ValidableJson>(
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<Transaction<'c, Postgres>> {
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),
}
}
backend/windmill-queue/src/schedule.rs
+2 -1
View File
@@ -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<RearmOutcome> {
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.
backend/windmill-store/src/resources.rs
+2 -1
View File
@@ -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
backend/windmill-trigger/src/listener.rs
+2 -1
View File
@@ -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#"
backend/windmill-worker/src/ai/tools.rs
+2 -1
View File
@@ -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)
backend/windmill-worker/src/ai/utils.rs
+2 -1
View File
@@ -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<String>,
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,
backend/windmill-worker/src/ai_executor.rs
+2 -1
View File
@@ -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!(
backend/windmill-worker/src/bun_executor.rs
+5 -4
View File
@@ -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,
backend/windmill-worker/src/dbt_executor.rs
+4 -3
View File
@@ -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
backend/windmill-worker/src/duckdb_executor.rs
+2 -1
View File
@@ -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<i64>,
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,
backend/windmill-worker/src/result_processor.rs
+2 -1
View File
@@ -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
backend/windmill-worker/src/worker_flow.rs
+8 -7
View File
@@ -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<Uuid>,
// 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)
backend/windmill-worker/src/worker_lockfiles.rs
+7 -6
View File
@@ -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?;