From 7ff804577c8f83fd36706dbaa00d2245c69ce413 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@windmill.dev>
Date: Mon, 24 Aug 2026 18:42:09 +0000
Subject: [PATCH] fix(postgres): roll back a transaction cancelled during BEGIN

`PgTransactionManager::begin` raises `transaction_depth` only after the BEGIN
round trip returns, but `start_rollback` -- which both its own `Rollback`
drop guard and `Transaction`'s drop guard call -- is a no-op while that depth
is zero. A future cancelled during the await therefore queues nothing, and the
session is left inside a transaction.

`Floating::return_to_pool` then validates the connection with `ping()`, which
for Postgres is a bare `wait_until_ready`: it drains the `ReadyForQuery` but
never inspects its transaction-status byte, so the connection is judged healthy
and handed to the next borrower. Their statements run inside the stale
transaction and hold its locks, and the first error turns the session into
`idle in transaction (aborted)`, after which every unrelated query on that
connection fails with 25P02 until `max_lifetime` recycles it.

Claim the depth before the round trip and unwind it if the BEGIN did not take,
so the drop guards have something to act on. The queued ROLLBACK is written to
the same buffer as the BEGIN and so is always flushed after it.

Closes #4393. Refs #2054, #2819, #3980.
---
 sqlx-postgres/src/transaction.rs |  9 +++++-
 tests/postgres/postgres.rs       | 51 ++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/sqlx-postgres/src/transaction.rs b/sqlx-postgres/src/transaction.rs
index 3f4122ea82..6ae0da53fc 100644
--- a/sqlx-postgres/src/transaction.rs
+++ b/sqlx-postgres/src/transaction.rs
@@ -27,11 +27,18 @@ impl TransactionManager for PgTransactionManager {
 
         let rollback = Rollback::new(conn);
         rollback.conn.queue_simple_query(statement.as_str())?;
+        // Claim the depth before the round trip, not after. `start_rollback` -- which both
+        // this guard and `Transaction`'s own drop guard call -- is a no-op while the depth is
+        // zero, so a future cancelled during the await below would otherwise leave the server
+        // in a transaction with nothing queued to end it. `Pool` then hands that connection
+        // to the next borrower, whose statements silently run inside it. Unwound just below
+        // if the BEGIN did not take.
+        rollback.conn.inner.transaction_depth += 1;
         rollback.conn.wait_until_ready().await?;
         if !rollback.conn.in_transaction() {
+            rollback.conn.inner.transaction_depth -= 1;
             return Err(Error::BeginFailed);
         }
-        rollback.conn.inner.transaction_depth += 1;
         rollback.defuse();
 
         Ok(())
diff --git a/tests/postgres/postgres.rs b/tests/postgres/postgres.rs
index 26f827b837..0e5f2829c9 100644
--- a/tests/postgres/postgres.rs
+++ b/tests/postgres/postgres.rs
@@ -2244,3 +2244,54 @@ async fn it_can_recover_from_copy_in_invalid_params() -> anyhow::Result<()> {
     )
     .await
 }
+
+// Regression: a future cancelled while `BEGIN`'s round trip is in flight used to leave the
+// session inside a transaction. `start_rollback` is a no-op while `transaction_depth` is
+// zero, and the depth was raised only after the await, so neither drop guard queued a
+// `ROLLBACK` -- and `return_to_pool` validates with a bare `wait_until_ready` that never
+// looks at the `ReadyForQuery` transaction-status byte, so the connection was handed to the
+// next borrower with the transaction still open.
+#[sqlx_macros::test]
+async fn it_rolls_back_a_transaction_cancelled_during_begin() -> anyhow::Result<()> {
+    let pool = PgPoolOptions::new()
+        .max_connections(1)
+        .min_connections(0)
+        .connect(&dotenvy::var("DATABASE_URL")?)
+        .await?;
+
+    let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()")
+        .fetch_one(&pool)
+        .await?;
+
+    // A plain `BEGIN` answers too quickly to cancel reliably; the sleep widens the same
+    // round trip so the cancellation lands inside it.
+    let cancelled = sqlx_core::rt::timeout(
+        Duration::from_millis(300),
+        pool.begin_with(AssertSqlSafe("BEGIN; SELECT pg_sleep(2);".to_string())),
+    )
+    .await;
+    assert!(cancelled.is_err(), "the begin should not have completed");
+
+    // Outlast the sleep: the queued `ROLLBACK` is only flushed once the abandoned statement
+    // has answered and the connection is on its way back to the pool.
+    sqlx_core::rt::sleep(Duration::from_millis(3500)).await;
+
+    let mut conn = new::<Postgres>().await?;
+    let state: Option<String> =
+        sqlx::query_scalar("SELECT state FROM pg_stat_activity WHERE pid = $1")
+            .bind(pid)
+            .fetch_optional(&mut conn)
+            .await?;
+
+    assert_eq!(
+        state.as_deref(),
+        Some("idle"),
+        "connection was returned to the pool still inside a transaction"
+    );
+
+    // and the pooled connection is still usable
+    let one: i32 = sqlx::query_scalar("SELECT 1").fetch_one(&pool).await?;
+    assert_eq!(one, 1);
+
+    Ok(())
+}