Improve the AnsiTransactionManager implementation to behave like that

one from diesel

Author: weiznich

src/lib.rs

@@ -140,31 +140,27 @@ where
     ///
     /// This function executes the provided closure `f` inside a database
     /// transaction. If there is already an open transaction for the current
-    /// connection savepoints will be used instead. The connection is commited if
+    /// connection savepoints will be used instead. The connection is committed if
     /// the closure returns `Ok(_)`, it will be rolled back if it returns `Err(_)`.
     /// For both cases the original result value will be returned from this function.
     ///
     /// If the transaction fails to commit due to a `SerializationFailure` or a
-    /// `ReadOnlyTransaction` a rollback will be attempted. In this case a
-    /// [`Error::CommitTransactionFailed`](crate::result::Error::CommitTransactionFailed)
-    /// error is returned, which contains details about the original error and
-    /// the success of the rollback attempt.
-    /// If the rollback failed the connection should be considered broken
+    /// `ReadOnlyTransaction` a rollback will be attempted.
+    /// If the rollback fails, the error will be returned in a
+    /// [`Error::RollbackErrorOnCommit`](crate::result::Error::RollbackErrorOnCommit),
+    /// from which you will be able to extract both the original commit error and
+    /// the rollback error.
+    /// In addition, the connection will be considered broken
     /// as it contains a uncommitted unabortable open transaction. Any further
     /// interaction with the transaction system will result in an returned error
-    /// in this cases.
+    /// in this case.
     ///
     /// If the closure returns an `Err(_)` and the rollback fails the function
-    /// will return a [`Error::RollbackError`](crate::result::Error::RollbackError)
-    /// wrapping the error generated by the rollback operation instead.
-    /// In this case the connection should be considered broken as it contains
-    /// an unabortable open transaction.
+    /// will return that rollback error directly, and the transaction manager will
+    /// be marked as broken as it contains a uncommitted unabortable open transaction.
     ///
     /// If a nested transaction fails to release the corresponding savepoint
-    /// a rollback will be attempted.  In this case a
-    /// [`Error::CommitTransactionFailed`](crate::result::Error::CommitTransactionFailed)
-    /// error is returned, which contains the original error and
-    /// details about the success of the rollback attempt.
+    /// the error will be returned directly.
     ///
     /// # Example
     ///
@@ -218,32 +214,28 @@ where
         E: From<diesel::result::Error> + Send,
         R: Send,
     {
-        Self::TransactionManager::begin_transaction(self).await?;
-        match callback(&mut *self).await {
-            Ok(value) => {
-                Self::TransactionManager::commit_transaction(self).await?;
-                Ok(value)
-            }
-            Err(user_error) => {
-                match Self::TransactionManager::rollback_transaction(self).await {
-                    Ok(()) => Err(user_error),
-                    Err(diesel::result::Error::BrokenTransactionManager) => {
-                        // In this case we are probably more interested by the
-                        // original error, which likely caused this
-                        Err(user_error)
-                    }
-                    Err(rollback_error) => Err(rollback_error.into()),
-                }
-            }
-        }
+        Self::TransactionManager::transaction(self, callback).await
     }
 
     /// Creates a transaction that will never be committed. This is useful for
     /// tests. Panics if called while inside of a transaction or
     /// if called with a connection containing a broken transaction
     async fn begin_test_transaction(&mut self) -> QueryResult<()> {
-        assert_eq!(Self::TransactionManager::get_transaction_depth(self), 0);
-        Self::TransactionManager::begin_transaction(self).await
+        use crate::transaction_manager::TransactionManagerStatus;
+
+        match Self::TransactionManager::transaction_manager_status_mut(self) {
+            TransactionManagerStatus::Valid(valid_status) => {
+                assert_eq!(None, valid_status.transaction_depth())
+            }
+            TransactionManagerStatus::InError => panic!("Transaction manager in error"),
+        };
+        Self::TransactionManager::begin_transaction(self).await?;
+        // set the test transaction flag
+        // to pervent that this connection gets droped in connection pools
+        // Tests commonly set the poolsize to 1 and use `begin_test_transaction`
+        // to prevent modifications to the schema
+        Self::TransactionManager::transaction_manager_status_mut(self).set_test_transaction_flag();
+        Ok(())
     }
 
     #[doc(hidden)]

src/mysql/mod.rs

@@ -126,6 +126,23 @@ impl AsyncConnection for AsyncMysqlConnection {
     }
 }
 
+#[inline(always)]
+fn update_transaction_manager_status<T>(
+    query_result: QueryResult<T>,
+    transaction_manager: &mut AnsiTransactionManager,
+) -> QueryResult<T> {
+    if let Err(diesel::result::Error::DatabaseError(
+        diesel::result::DatabaseErrorKind::SerializationFailure,
+        _,
+    )) = query_result
+    {
+        transaction_manager
+            .status
+            .set_top_level_transaction_requires_rollback()
+    }
+    query_result
+}
+
 #[async_trait::async_trait]
 impl PrepareCallback<Statement, MysqlType> for &'_ mut mysql_async::Conn {
     async fn prepare(
@@ -194,6 +211,7 @@ impl AsyncMysqlConnection {
             ref mut conn,
             ref mut stmt_cache,
             ref mut last_stmt,
+            ref mut transaction_manager,
             ..
         } = self;
 
@@ -209,7 +227,7 @@ impl AsyncMysqlConnection {
                 MaybeCached::Cached(s) => s,
                 _ => unreachable!("We've opted into breaking diesel changes and want to know if things break because someone added a new variant here")
             };
-            callback(conn, stmt, ToSqlHelper{metadata, binds}).await
+            update_transaction_manager_status(callback(conn, stmt, ToSqlHelper{metadata, binds}).await, transaction_manager)
         }).boxed()
     }
 }

src/pg/mod.rs

@@ -4,7 +4,6 @@ use self::serialize::ToSqlHelper;
 use crate::stmt_cache::{PrepareCallback, StmtCache};
 use crate::{
     AnsiTransactionManager, AsyncConnection, AsyncConnectionGatWorkaround, SimpleAsyncConnection,
-    TransactionManager,
 };
 use diesel::connection::statement_cache::PrepareForCache;
 use diesel::pg::{
@@ -94,7 +93,7 @@ mod transaction_builder;
 pub struct AsyncPgConnection {
     conn: Arc<tokio_postgres::Client>,
     stmt_cache: Arc<Mutex<StmtCache<diesel::pg::Pg, Statement>>>,
-    transaction_state: AnsiTransactionManager,
+    transaction_state: Arc<Mutex<AnsiTransactionManager>>,
     metadata_cache: Arc<Mutex<Option<PgMetadataCache>>>,
 }
 
@@ -152,11 +151,13 @@ impl AsyncConnection for AsyncPgConnection {
         let conn = self.conn.clone();
         let stmt_cache = self.stmt_cache.clone();
         let metadata_cache = self.metadata_cache.clone();
+        let tm = self.transaction_state.clone();
         let query = source.as_query();
         Self::with_prepared_statement(
             conn,
             stmt_cache,
             metadata_cache,
+            tm,
             query,
             |conn, stmt, binds| async move {
                 let res = conn.query_raw(&stmt, binds).await.map_err(ErrorHelper)?;
@@ -181,6 +182,7 @@ impl AsyncConnection for AsyncPgConnection {
             self.conn.clone(),
             self.stmt_cache.clone(),
             self.metadata_cache.clone(),
+            self.transaction_state.clone(),
             source,
             |conn, stmt, binds| async move {
                 let binds = binds
@@ -197,42 +199,33 @@ impl AsyncConnection for AsyncPgConnection {
         .boxed()
     }
 
-    fn transaction_state(
-        &mut self,
-    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
-        &mut self.transaction_state
-    }
-
-    async fn transaction<R, E, F>(&mut self, callback: F) -> Result<R, E>
-    where
-        F: FnOnce(&mut Self) -> futures::future::BoxFuture<Result<R, E>> + Send,
-        E: From<diesel::result::Error> + Send,
-        R: Send,
-    {
-        Self::TransactionManager::begin_transaction(self).await?;
-        match callback(&mut *self).await {
-            Ok(value) => {
-                Self::TransactionManager::commit_transaction(self).await?;
-                Ok(value)
-            }
-            Err(user_error) => {
-                match Self::TransactionManager::rollback_transaction(self).await {
-                    Ok(()) => Err(user_error),
-                    Err(diesel::result::Error::BrokenTransactionManager) => {
-                        // In this case we are probably more interested by the
-                        // original error, which likely caused this
-                        Err(user_error)
-                    }
-                    Err(rollback_error) => Err(rollback_error.into()),
-                }
-            }
+    fn transaction_state(&mut self) -> &mut AnsiTransactionManager {
+        // there should be no other pending future when this is called
+        // that means there is only one instance of this arc and
+        // we can simply access the inner data
+        if let Some(tm) = Arc::get_mut(&mut self.transaction_state) {
+            tm.get_mut()
+        } else {
+            panic!("Cannot access shared transaction state")
         }
     }
+}
 
-    async fn begin_test_transaction(&mut self) -> QueryResult<()> {
-        assert_eq!(Self::TransactionManager::get_transaction_depth(self), 0);
-        Self::TransactionManager::begin_transaction(self).await
+#[inline(always)]
+fn update_transaction_manager_status<T>(
+    query_result: QueryResult<T>,
+    transaction_manager: &mut AnsiTransactionManager,
+) -> QueryResult<T> {
+    if let Err(diesel::result::Error::DatabaseError(
+        diesel::result::DatabaseErrorKind::SerializationFailure,
+        _,
+    )) = query_result
+    {
+        transaction_manager
+            .status
+            .set_top_level_transaction_requires_rollback()
     }
+    query_result
 }
 
 #[async_trait::async_trait]
@@ -308,7 +301,7 @@ impl AsyncPgConnection {
         let mut conn = Self {
             conn: Arc::new(conn),
             stmt_cache: Arc::new(Mutex::new(StmtCache::new())),
-            transaction_state: AnsiTransactionManager::default(),
+            transaction_state: Arc::new(Mutex::new(AnsiTransactionManager::default())),
             metadata_cache: Arc::new(Mutex::new(Some(PgMetadataCache::new()))),
         };
         conn.set_config_options()
@@ -333,6 +326,7 @@ impl AsyncPgConnection {
         raw_connection: Arc<tokio_postgres::Client>,
         stmt_cache: Arc<Mutex<StmtCache<diesel::pg::Pg, Statement>>>,
         metadata_cache: Arc<Mutex<Option<PgMetadataCache>>>,
+        tm: Arc<Mutex<AnsiTransactionManager>>,
         query: T,
         callback: impl FnOnce(Arc<tokio_postgres::Client>, Statement, Vec<ToSqlHelper>) -> F,
     ) -> QueryResult<R>
@@ -396,15 +390,15 @@ impl AsyncPgConnection {
                     }
                 } else {
                     // bubble up any error as soon as we have done all lookups
-                    let _ = res?;
+                    res?;
                     break;
                 }
             }
         }
 
         let stmt = {
             let mut stmt_cache = stmt_cache.lock().await;
-            let stmt = stmt_cache
+            stmt_cache
                 .cached_prepared_statement(
                     query,
                     &bind_collector.metadata,
@@ -413,8 +407,7 @@ impl AsyncPgConnection {
                 )
                 .await?
                 .0
-                .clone();
-            stmt
+                .clone()
         };
 
         let binds = bind_collector
@@ -423,7 +416,9 @@ impl AsyncPgConnection {
             .zip(bind_collector.binds)
             .map(|(meta, bind)| ToSqlHelper(meta, bind))
             .collect::<Vec<_>>();
-        callback(raw_connection, stmt.clone(), binds).await
+        let res = callback(raw_connection, stmt.clone(), binds).await;
+        let mut tm = tm.lock().await;
+        update_transaction_manager_status(res, &mut *tm)
     }
 }
 

src/pooled_connection/mod.rs

@@ -173,8 +173,10 @@ where
         TM::commit_transaction(&mut **conn).await
     }
 
-    fn get_transaction_depth(conn: &mut C) -> u32 {
-        TM::get_transaction_depth(&mut **conn)
+    fn transaction_manager_status_mut(
+        conn: &mut C,
+    ) -> &mut crate::transaction_manager::TransactionManagerStatus {
+        TM::transaction_manager_status_mut(&mut **conn)
     }
 }