From f93bbe87bec6c7ff480b783e9a6400063fde305d Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <oon@conradludgate.com>
Date: Tue, 17 Oct 2023 03:20:48 +0100
Subject: [PATCH] postgres begin cancel safe (#2819)

---
 sqlx-postgres/src/transaction.rs | 35 ++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/sqlx-postgres/src/transaction.rs b/sqlx-postgres/src/transaction.rs
index 045dda70..02028624 100644
--- a/sqlx-postgres/src/transaction.rs
+++ b/sqlx-postgres/src/transaction.rs
@@ -15,10 +15,12 @@ impl TransactionManager for PgTransactionManager {
 
     fn begin(conn: &mut PgConnection) -> BoxFuture<'_, Result<(), Error>> {
         Box::pin(async move {
-            conn.execute(&*begin_ansi_transaction_sql(conn.transaction_depth))
-                .await?;
-
-            conn.transaction_depth += 1;
+            let rollback = Rollback::new(conn);
+            let query = begin_ansi_transaction_sql(rollback.conn.transaction_depth);
+            rollback.conn.queue_simple_query(&query);
+            rollback.conn.transaction_depth += 1;
+            rollback.conn.wait_until_ready().await?;
+            rollback.defuse();
 
             Ok(())
         })
@@ -58,3 +60,28 @@ impl TransactionManager for PgTransactionManager {
         }
     }
 }
+
+struct Rollback<'c> {
+    conn: &'c mut PgConnection,
+    defuse: bool,
+}
+
+impl Drop for Rollback<'_> {
+    fn drop(&mut self) {
+        if !self.defuse {
+            PgTransactionManager::start_rollback(self.conn)
+        }
+    }
+}
+
+impl<'c> Rollback<'c> {
+    fn new(conn: &'c mut PgConnection) -> Self {
+        Self {
+            conn,
+            defuse: false,
+        }
+    }
+    fn defuse(mut self) {
+        self.defuse = true;
+    }
+}