From f94c33d1eaf7978b39d3ad41df32b185d77d4fde Mon Sep 17 00:00:00 2001
From: Chris Tsang <tyt2y7@gmail.com>
Date: Mon, 11 Oct 2021 12:11:40 +0800
Subject: [PATCH] Explicit COMMIT and ROLLBACK in Mock

---
 src/database/mock.rs               | 23 ++++++++++++++++++++---
 src/database/stream/transaction.rs |  2 ++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/database/mock.rs b/src/database/mock.rs
index 8bde4725..bdab864c 100644
--- a/src/database/mock.rs
+++ b/src/database/mock.rs
@@ -117,7 +117,8 @@ impl MockDatabaseTrait for MockDatabase {
 
     fn commit(&mut self) {
         if self.transaction.is_some() {
-            let transaction = self.transaction.take().unwrap();
+            let mut transaction = self.transaction.take().unwrap();
+            transaction.push(Statement::from_string(self.db_backend, "COMMIT".to_owned()));
             self.transaction_log.push(transaction.into_transaction());
         } else {
             panic!("There is no open transaction to commit");
@@ -126,7 +127,12 @@ impl MockDatabaseTrait for MockDatabase {
 
     fn rollback(&mut self) {
         if self.transaction.is_some() {
-            self.transaction = None;
+            let mut transaction = self.transaction.take().unwrap();
+            transaction.push(Statement::from_string(
+                self.db_backend,
+                "ROLLBACK".to_owned(),
+            ));
+            self.transaction_log.push(transaction.into_transaction());
         } else {
             panic!("There is no open transaction to rollback");
         }
@@ -283,6 +289,7 @@ mod tests {
                         r#"SELECT "fruit"."id", "fruit"."name", "fruit"."cake_id" FROM "fruit""#,
                         vec![]
                     ),
+                    Statement::from_string(DbBackend::Postgres, "COMMIT".to_owned()),
                 ]),
                 Transaction::from_sql_and_values(
                     DbBackend::Postgres,
@@ -313,7 +320,17 @@ mod tests {
             _ => panic!(),
         }
 
-        assert_eq!(db.into_transaction_log(), vec![]);
+        assert_eq!(
+            db.into_transaction_log(),
+            vec![Transaction::many(vec![
+                Statement::from_sql_and_values(
+                    DbBackend::Postgres,
+                    r#"SELECT "cake"."id", "cake"."name" FROM "cake" LIMIT $1"#,
+                    vec![1u64.into()]
+                ),
+                Statement::from_string(DbBackend::Postgres, "ROLLBACK".to_owned()),
+            ]),]
+        );
     }
 
     #[smol_potat::test]
diff --git a/src/database/stream/transaction.rs b/src/database/stream/transaction.rs
index 651f3d11..2dddc59c 100644
--- a/src/database/stream/transaction.rs
+++ b/src/database/stream/transaction.rs
@@ -12,6 +12,8 @@ use futures::lock::MutexGuard;
 use crate::{DbErr, InnerConnection, QueryResult, Statement};
 
 #[ouroboros::self_referencing]
+/// `TransactionStream` cannot be used in a `transaction` closure as it does not impl `Send`.
+/// It seems to be a Rust limitation right now, and solution to work around this deemed to be extremely hard.
 pub struct TransactionStream<'a> {
     stmt: Statement,
     conn: MutexGuard<'a, InnerConnection>,