Merge pull request #60 from KatsuoRyuu/feature/test_transaction

Adding test transaction to the trait AsyncConnection

Author: weiznich

src/lib.rs

@@ -66,14 +66,17 @@
 //! ```
 
 #![warn(missing_docs)]
+
 use diesel::backend::Backend;
 use diesel::query_builder::{AsQuery, QueryFragment, QueryId};
+use diesel::result::Error;
 use diesel::row::Row;
 use diesel::{ConnectionResult, QueryResult};
 use futures_util::{Future, Stream};
+use std::fmt::Debug;
 
 pub use scoped_futures;
-use scoped_futures::ScopedBoxFuture;
+use scoped_futures::{ScopedBoxFuture, ScopedFutureExt};
 
 #[cfg(feature = "mysql")]
 mod mysql;
@@ -254,6 +257,65 @@ pub trait AsyncConnection: SimpleAsyncConnection + Sized + Send {
         Ok(())
     }
 
+    /// Executes the given function inside a transaction, but does not commit
+    /// it. Panics if the given function returns an error.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # include!("doctest_setup.rs");
+    /// use diesel::result::Error;
+    /// use scoped_futures::ScopedFutureExt;
+    ///
+    /// # #[tokio::main(flavor = "current_thread")]
+    /// # async fn main() {
+    /// #     run_test().await.unwrap();
+    /// # }
+    /// #
+    /// # async fn run_test() -> QueryResult<()> {
+    /// #     use schema::users::dsl::*;
+    /// #     let conn = &mut establish_connection().await;
+    /// conn.test_transaction::<_, Error, _>(|conn| async move {
+    ///     diesel::insert_into(users)
+    ///         .values(name.eq("Ruby"))
+    ///         .execute(conn)
+    ///         .await?;
+    ///
+    ///     let all_names = users.select(name).load::<String>(conn).await?;
+    ///     assert_eq!(vec!["Sean", "Tess", "Ruby"], all_names);
+    ///
+    ///     Ok(())
+    /// }.scope_boxed()).await;
+    ///
+    /// // Even though we returned `Ok`, the transaction wasn't committed.
+    /// let all_names = users.select(name).load::<String>(conn).await?;
+    /// assert_eq!(vec!["Sean", "Tess"], all_names);
+    /// #     Ok(())
+    /// # }
+    /// ```
+    async fn test_transaction<'a, R, E, F>(&'a mut self, f: F) -> R
+    where
+        F: for<'r> FnOnce(&'r mut Self) -> ScopedBoxFuture<'a, 'r, Result<R, E>> + Send + 'a,
+        E: Debug + Send + 'a,
+        R: Send + 'a,
+        Self: 'a,
+    {
+        use futures_util::TryFutureExt;
+
+        let mut user_result = None;
+        let _ = self
+            .transaction::<R, _, _>(|c| {
+                f(c).map_err(|_| Error::RollbackTransaction)
+                    .and_then(|r| {
+                        user_result = Some(r);
+                        futures_util::future::ready(Err(Error::RollbackTransaction))
+                    })
+                    .scope_boxed()
+            })
+            .await;
+        user_result.expect("Transaction did not succeed")
+    }
+
     #[doc(hidden)]
     fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
     where

src/pg/mod.rs

@@ -123,7 +123,7 @@ impl AsyncConnection for AsyncPgConnection {
             .map_err(ErrorHelper)?;
         tokio::spawn(async move {
             if let Err(e) = connection.await {
-                eprintln!("connection error: {}", e);
+                eprintln!("connection error: {e}");
             }
         });
         Self::try_from(client).await

src/pooled_connection/mod.rs

@@ -39,14 +39,17 @@ impl fmt::Display for PoolError {
 
 impl std::error::Error for PoolError {}
 
+type SetupCallback<C> =
+    Box<dyn Fn(&str) -> future::BoxFuture<diesel::ConnectionResult<C>> + Send + Sync>;
+
 /// An connection manager for use with diesel-async.
 ///
 /// See the concrete pool implementations for examples:
 /// * [deadpool](self::deadpool)
 /// * [bb8](self::bb8)
 /// * [mobc](self::mobc)
 pub struct AsyncDieselConnectionManager<C> {
-    setup: Box<dyn Fn(&str) -> future::BoxFuture<diesel::ConnectionResult<C>> + Send + Sync>,
+    setup: SetupCallback<C>,
     connection_url: String,
 }
 

src/transaction_manager.rs

@@ -292,7 +292,7 @@ where
         let start_transaction_sql = match transaction_state.transaction_depth() {
             None => Cow::from("BEGIN"),
             Some(transaction_depth) => {
-                Cow::from(format!("SAVEPOINT diesel_savepoint_{}", transaction_depth))
+                Cow::from(format!("SAVEPOINT diesel_savepoint_{transaction_depth}"))
             }
         };
         conn.batch_execute(&start_transaction_sql).await?;

tests/type_check.rs

@@ -161,8 +161,8 @@ async fn test_timestamp() {
     type_check::<_, sql_types::Timestamp>(
         conn,
         chrono::NaiveDateTime::new(
-            chrono::NaiveDate::from_ymd(2021, 09, 27),
-            chrono::NaiveTime::from_hms_milli(17, 44, 23, 0),
+            chrono::NaiveDate::from_ymd_opt(2021, 09, 27).unwrap(),
+            chrono::NaiveTime::from_hms_milli_opt(17, 44, 23, 0).unwrap(),
         ),
     )
     .await;
@@ -171,13 +171,18 @@ async fn test_timestamp() {
 #[tokio::test]
 async fn test_date() {
     let conn = &mut connection().await;
-    type_check::<_, sql_types::Date>(conn, chrono::NaiveDate::from_ymd(2021, 09, 27)).await;
+    type_check::<_, sql_types::Date>(conn, chrono::NaiveDate::from_ymd_opt(2021, 09, 27).unwrap())
+        .await;
 }
 
 #[tokio::test]
 async fn test_time() {
     let conn = &mut connection().await;
-    type_check::<_, sql_types::Time>(conn, chrono::NaiveTime::from_hms_milli(17, 44, 23, 0)).await;
+    type_check::<_, sql_types::Time>(
+        conn,
+        chrono::NaiveTime::from_hms_milli_opt(17, 44, 23, 0).unwrap(),
+    )
+    .await;
 }
 
 #[cfg(feature = "mysql")]
@@ -187,8 +192,8 @@ async fn test_datetime() {
     type_check::<_, sql_types::Datetime>(
         conn,
         chrono::NaiveDateTime::new(
-            chrono::NaiveDate::from_ymd(2021, 09, 30),
-            chrono::NaiveTime::from_hms_milli(12, 06, 42, 0),
+            chrono::NaiveDate::from_ymd_opt(2021, 09, 30).unwrap(),
+            chrono::NaiveTime::from_hms_milli_opt(12, 06, 42, 0).unwrap(),
         ),
     )
     .await;