transaction: document, test, and fix operation of nested transactions

Author: mehcode

sqlx-core/src/transaction.rs

@@ -9,6 +9,7 @@ use crate::describe::Describe;
 use crate::executor::{Execute, Executor, RefExecutor};
 use crate::runtime::spawn;
 
+/// Represents a database transaction.
 // Transaction<PoolConnection<PgConnection>>
 // Transaction<PgConnection>
 pub struct Transaction<T>
@@ -38,10 +39,14 @@ where
  })
  }
 
- pub async fn begin(mut self) -> crate::Result<Transaction<T>> {
- Transaction::new(self.depth, self.inner.take().expect(ERR_FINALIZED)).await
+ /// Creates a new save point in the current transaction and returns
+ /// a new `Transaction` object to manage its scope.
+ pub async fn begin(self) -> crate::Result<Transaction<Transaction<T>>> {
+ Transaction::new(self.depth, self).await
  }
 
+ /// Commits the current transaction or save point.
+ /// Returns the inner connection or transaction.
  pub async fn commit(mut self) -> crate::Result<T> {
  let mut inner = self.inner.take().expect(ERR_FINALIZED);
  let depth = self.depth;
@@ -57,6 +62,8 @@ where
  Ok(inner)
  }
 
+ /// Rollback the current transaction or save point.
+ /// Returns the inner connection or transaction.
  pub async fn rollback(mut self) -> crate::Result<T> {
  let mut inner = self.inner.take().expect(ERR_FINALIZED);
  let depth = self.depth;
@@ -95,15 +102,49 @@ where
  }
 }
 
-impl<'c, DB, T> Executor for &'c mut Transaction<T>
+impl<T> Connection for Transaction<T>
+where
+ T: Connection,
+{
+ // Close is equivalent to
+ fn close(mut self) -> BoxFuture<'static, crate::Result<()>> {
+ Box::pin(async move {
+ let mut inner = self.inner.take().expect(ERR_FINALIZED);
+
+ if self.depth == 1 {
+ // This is the root transaction, call rollback
+ let res = inner.execute("ROLLBACK").await;
+
+ // No matter the result of the above, call close
+ let _ = inner.close().await;
+
+ // Now raise the error if there was one
+ res?;
+ } else {
+ // This is not the root transaction, forward to a nested
+ // transaction (to eventually call rollback)
+ inner.close().await?
+ }
+
+ Ok(())
+ })
+ }
+
+ #[inline]
+ fn ping(&mut self) -> BoxFuture<'_, crate::Result<()>> {
+ self.deref_mut().ping()
+ }
+}
+
+impl<DB, T> Executor for Transaction<T>
 where
  DB: Database,
  T: Connection<Database = DB>,
 {
  type Database = T::Database;
 
- fn execute<'e, 'q: 'e, 't: 'e, E: 'e>(
- &'t mut self,
+ fn execute<'e, 'q: 'e, 'c: 'e, E: 'e>(
+ &'c mut self,
  query: E,
  ) -> BoxFuture<'e, crate::Result<u64>>
  where
@@ -112,7 +153,7 @@ where
  (**self).execute(query)
  }
 
- fn fetch<'q, 'e, E>(&'e mut self, query: E) -> <Self::Database as HasCursor<'e, 'q>>::Cursor
+ fn fetch<'e, 'q, E>(&'e mut self, query: E) -> <Self::Database as HasCursor<'e, 'q>>::Cursor
  where
  E: Execute<'q, Self::Database>,
  {
@@ -151,16 +192,9 @@ where
 {
  fn drop(&mut self) {
  if self.depth > 0 {
- if let Some(mut inner) = self.inner.take() {
+ if let Some(inner) = self.inner.take() {
  spawn(async move {
- let res = inner.execute("ROLLBACK").await;
-
- // If the rollback failed we need to close the inner connection
- if res.is_err() {
- // This will explicitly forget the connection so it will not
- // return to the pool
- let _ = inner.close().await;
- }
+ let _ = inner.close().await;
  });
  }
  }

tests/postgres.rs

@@ -153,6 +153,99 @@ async fn it_can_work_with_transactions() -> anyhow::Result<()> {
  Ok(())
 }
 
+#[cfg_attr(feature = "runtime-async-std", async_std::test)]
+#[cfg_attr(feature = "runtime-tokio", tokio::test)]
+async fn it_can_work_with_nested_transactions() -> anyhow::Result<()> {
+ let mut conn = new::<Postgres>().await?;
+
+ conn.execute("CREATE TABLE IF NOT EXISTS _sqlx_users_2523 (id INTEGER PRIMARY KEY)")
+ .await?;
+
+ conn.execute("TRUNCATE _sqlx_users_2523").await?;
+
+ // begin
+ let mut tx = conn.begin().await?;
+
+ // insert a user
+ sqlx::query("INSERT INTO _sqlx_users_2523 (id) VALUES ($1)")
+ .bind(50_i32)
+ .execute(&mut tx)
+ .await?;
+
+ // begin once more
+ let mut tx = tx.begin().await?;
+
+ // insert another user
+ sqlx::query("INSERT INTO _sqlx_users_2523 (id) VALUES ($1)")
+ .bind(10_i32)
+ .execute(&mut tx)
+ .await?;
+
+ // never mind, rollback
+ let mut tx = tx.rollback().await?;
+
+ // did we really?
+ let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM _sqlx_users_2523")
+ .fetch_one(&mut tx)
+ .await?;
+
+ assert_eq!(count, 1);
+
+ // actually, commit
+ let mut conn = tx.commit().await?;
+
+ // did we really?
+ let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM _sqlx_users_2523")
+ .fetch_one(&mut conn)
+ .await?;
+
+ assert_eq!(count, 1);
+
+ Ok(())
+}
+
+#[cfg_attr(feature = "runtime-async-std", async_std::test)]
+#[cfg_attr(feature = "runtime-tokio", tokio::test)]
+async fn it_can_rollback_nested_transactions() -> anyhow::Result<()> {
+ let mut conn = new::<Postgres>().await?;
+
+ conn.execute("CREATE TABLE IF NOT EXISTS _sqlx_users_512412 (id INTEGER PRIMARY KEY)")
+ .await?;
+
+ conn.execute("TRUNCATE _sqlx_users_512412").await?;
+
+ // begin
+ let mut tx = conn.begin().await?;
+
+ // insert a user
+ sqlx::query("INSERT INTO _sqlx_users_512412 (id) VALUES ($1)")
+ .bind(50_i32)
+ .execute(&mut tx)
+ .await?;
+
+ // begin once more
+ let mut tx = tx.begin().await?;
+
+ // insert another user
+ sqlx::query("INSERT INTO _sqlx_users_512412 (id) VALUES ($1)")
+ .bind(10_i32)
+ .execute(&mut tx)
+ .await?;
+
+ // stop the phone, drop the entire transaction
+ tx.close().await?;
+
+ // did we really?
+ let mut conn = new::<Postgres>().await?;
+ let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM _sqlx_users_512412")
+ .fetch_one(&mut conn)
+ .await?;
+
+ assert_eq!(count, 0);
+
+ Ok(())
+}
+
 // run with `cargo test --features postgres -- --ignored --nocapture pool_smoke_test`
 #[ignore]
 #[cfg_attr(feature = "runtime-async-std", async_std::test)]