Skip to content

Commit

Permalink
postgres begin cancel safe (#2819)
Browse files Browse the repository at this point in the history
  • Loading branch information
conradludgate authored Oct 17, 2023
1 parent 58cb18a commit f93bbe8
Showing 1 changed file with 31 additions and 4 deletions.
35 changes: 31 additions & 4 deletions sqlx-postgres/src/transaction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(())
})
Expand Down Expand Up @@ -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;
}
}

0 comments on commit f93bbe8

Please sign in to comment.