From 50c7f229e9123a55531296d4b22d166b77ff7ca3 Mon Sep 17 00:00:00 2001
From: "He, Wanchen" <hwc14@qq.com>
Date: Mon, 11 Nov 2024 12:12:34 +0800
Subject: [PATCH] Add coro commit api for transaction.

---
 orm_lib/inc/drogon/orm/DbClient.h |  5 ++++
 orm_lib/src/TransactionImpl.cc    | 25 ++++++++++++++--
 orm_lib/src/TransactionImpl.h     |  6 ++++
 orm_lib/tests/db_test.cc          | 49 +++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/orm_lib/inc/drogon/orm/DbClient.h b/orm_lib/inc/drogon/orm/DbClient.h
index 89b3f12390..1291a86052 100644
--- a/orm_lib/inc/drogon/orm/DbClient.h
+++ b/orm_lib/inc/drogon/orm/DbClient.h
@@ -371,6 +371,11 @@ class Transaction : public DbClient
     void closeAll() override
     {
     }
+
+#ifdef __cpp_impl_coroutine
+    virtual void setAutoCommit(bool) = 0;
+    virtual drogon::Task<> commitCoro() = 0;
+#endif
 };
 
 #ifdef __cpp_impl_coroutine
diff --git a/orm_lib/src/TransactionImpl.cc b/orm_lib/src/TransactionImpl.cc
index 0f55142400..bfefd840b3 100644
--- a/orm_lib/src/TransactionImpl.cc
+++ b/orm_lib/src/TransactionImpl.cc
@@ -40,6 +40,7 @@ TransactionImpl::~TransactionImpl()
     {
         auto loop = connectionPtr_->loop();
         loop->queueInLoop([conn = connectionPtr_,
+                           autoCommit = autoCommit_,
                            ucb = std::move(usedUpCallback_),
                            commitCb = std::move(commitCallback_)]() {
             conn->setIdleCallback([ucb = std::move(ucb)]() {
@@ -47,16 +48,16 @@ TransactionImpl::~TransactionImpl()
                     ucb();
             });
             conn->execSql(
-                "commit",
+                autoCommit ? "commit" : "rollback",
                 0,
                 {},
                 {},
                 {},
-                [commitCb](const Result &) {
+                [commitCb, autoCommit](const Result &) {
                     LOG_TRACE << "Transaction committed!";
                     if (commitCb)
                     {
-                        commitCb(true);
+                        commitCb(autoCommit);
                     }
                 },
                 [commitCb](const std::exception_ptr &ePtr) {
@@ -389,3 +390,21 @@ void TransactionImpl::execSqlInLoopWithTimeout(
     }
     timeoutFlagPtr->runTimer();
 }
+
+#ifdef __cpp_impl_coroutine
+void TransactionImpl::setAutoCommit(bool autoCommit)
+{
+    autoCommit_ = autoCommit;
+}
+
+drogon::Task<> TransactionImpl::commitCoro()
+{
+    co_await execSqlCoro("commit");
+    isCommitedOrRolledback_ = true;
+    if (commitCallback_)
+    {
+        commitCallback_(true);
+    }
+    co_return;
+}
+#endif
diff --git a/orm_lib/src/TransactionImpl.h b/orm_lib/src/TransactionImpl.h
index 0be310f5c0..8b57e63447 100644
--- a/orm_lib/src/TransactionImpl.h
+++ b/orm_lib/src/TransactionImpl.h
@@ -50,6 +50,11 @@ class TransactionImpl : public Transaction,
         timeout_ = timeout;
     }
 
+#ifdef __cpp_impl_coroutine
+    void setAutoCommit(bool b) override;
+    Task<> commitCoro() override;
+#endif
+
   private:
     DbConnectionPtr connectionPtr_;
 
@@ -127,6 +132,7 @@ class TransactionImpl : public Transaction,
 
     std::function<void()> usedUpCallback_;
     bool isCommitedOrRolledback_{false};
+    bool autoCommit_{true};
     bool isWorking_{false};
     void execNewTask();
 
diff --git a/orm_lib/tests/db_test.cc b/orm_lib/tests/db_test.cc
index 086a8cca25..aed3d02fd5 100644
--- a/orm_lib/tests/db_test.cc
+++ b/orm_lib/tests/db_test.cc
@@ -1077,6 +1077,55 @@ DROGON_TEST(PostgreTest)
                 "what():",
                 e.base().what());
         }
+        // transaction auto commit
+        try
+        {
+            {
+                auto trans = co_await clientPtr->newTransactionCoro();
+                trans->setAutoCommit(false);
+                co_await trans->execSqlCoro(
+                    "update users set user_name = $1 where user_id = $2",
+                    "should not be this",
+                    "pg");
+                MANDATE(true);
+            }
+            auto result = co_await clientPtr->execSqlCoro(
+                "select user_name from users where user_id = $1", "pg");
+            MANDATE(result[0][0].as<std::string>() == "postgres");
+            LOG_INFO << "setAutoCommit success";
+        }
+        catch (const DrogonDbException &e)
+        {
+            FAULT(
+                "postgresql - DbClient coroutine transaction interface(1) "
+                "what():",
+                e.base().what());
+        }
+        // transaction commit
+        try
+        {
+            {
+                auto trans = co_await clientPtr->newTransactionCoro();
+                trans->setAutoCommit(false);
+                co_await trans->execSqlCoro(
+                    "update users set signature = $1 where user_id = $2",
+                    "commitCoro success",
+                    "pg");
+                MANDATE(true);
+                co_await trans->commitCoro();
+            }
+            auto result = co_await clientPtr->execSqlCoro(
+                "select signature from users where user_id = $1", "pg");
+            MANDATE(result[0][0].as<std::string>() == "commitCoro success");
+            LOG_INFO << "commitCoro success";
+        }
+        catch (const DrogonDbException &e)
+        {
+            FAULT(
+                "postgresql - DbClient coroutine transaction interface(2) "
+                "what():",
+                e.base().what());
+        }
     };
     drogon::sync_wait(coro_test());