From 8866a192edd1d6cf4915640fd78883f04b7628f2 Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Sun, 1 May 2022 00:07:14 +0800
Subject: [PATCH] Works with transaction, close #27, #15

---
 README.md          | 24 +++++++++++++++++++-----
 clauses.go         | 15 ++++++++++++---
 dbresolver_test.go | 38 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 68 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 59e0140..162f438 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,7 @@ DBResolver adds multiple databases support to GORM, the following features are s
 * Manual connection switching
 * Sources/Replicas load balancing
 * Works for RAW SQL
+* Transaction
 
 ## Quick Start
 
@@ -36,13 +37,9 @@ DB.Use(dbresolver.Register(dbresolver.Config{
 }, "orders", &Product{}, "secondary"))
 ```
 
-### Transaction
-
-When using transaction, DBResolver will use the transaction and won't switch to sources/replicas
-
 ### Automatic connection switching
 
-DBResolver will automatically switch connection based on the working table/struct
+DBResolver will automatically switch connections based on the working table/struct
 
 For RAW SQL, DBResolver will extract the table name from the SQL to match the resolver, and will use `sources` unless the SQL begins with `SELECT`, for example:
 
@@ -85,6 +82,23 @@ DB.Clauses(dbresolver.Use("secondary")).First(&user)
 DB.Clauses(dbresolver.Use("secondary"), dbresolver.Write).First(&user)
 ```
 
+### Transaction
+
+When using transaction, DBResolver will keep using the transaction and won't switch to sources/replicas based on configuration
+
+But you can specifies which DB to use before starting a transaction, for example:
+
+```go
+// Start transaction based on default replicas db
+tx := DB.Clauses(dbresolver.Read).Begin()
+
+// Start transaction based on default sources db
+tx := DB.Clauses(dbresolver.Read).Begin()
+
+// Start transaction based on `secondary`'s sources
+tx := DB.Clauses(dbresolver.Use("secondary"), dbresolver.Write).Begin()
+```
+
 ### Load Balancing
 
 GORM supports load balancing sources/replicas based on policy, the policy is an interface implements following interface:
diff --git a/clauses.go b/clauses.go
index 3a81da0..3a205a8 100644
--- a/clauses.go
+++ b/clauses.go
@@ -8,16 +8,24 @@ import (
 // Operation specifies dbresolver mode
 type Operation string
 
-const writeName = "gorm:db_resolver:write"
+const (
+	writeName = "gorm:db_resolver:write"
+	readName  = "gorm:db_resolver:read"
+)
 
 // ModifyStatement modify operation mode
 func (op Operation) ModifyStatement(stmt *gorm.Statement) {
-	optName := "gorm:db_resolver:read"
+	var optName string
 	if op == Write {
 		optName = writeName
+	} else if op == Read {
+		optName = readName
 	}
 
-	stmt.Clauses[optName] = clause.Clause{}
+	if optName != "" {
+		stmt.Clauses[optName] = clause.Clause{}
+		stmt.DB.Callback().Query().Get("gorm:db_resolver")(stmt.DB)
+	}
 }
 
 // Build implements clause.Expression interface
@@ -38,6 +46,7 @@ const usingName = "gorm:db_resolver:using"
 // ModifyStatement modify operation mode
 func (u using) ModifyStatement(stmt *gorm.Statement) {
 	stmt.Clauses[usingName] = clause.Clause{Expression: u}
+	stmt.DB.Callback().Query().Get("gorm:db_resolver")(stmt.DB)
 }
 
 // Build implements clause.Expression interface
diff --git a/dbresolver_test.go b/dbresolver_test.go
index 47e3987..34e22ba 100644
--- a/dbresolver_test.go
+++ b/dbresolver_test.go
@@ -65,8 +65,44 @@ func TestDBResolver(t *testing.T) {
 		}
 
 		for j := 0; j < 20; j++ {
-			// test query
 			var order Order
+			// test transaction
+			tx := DB.Begin()
+			tx.Find(&order)
+			if order.OrderNo != "9911" {
+				t.Fatalf("idx: %v: order should comes from default db, but got order %v", j, order.OrderNo)
+			}
+			tx.Rollback()
+
+			tx = DB.Clauses(dbresolver.Read).Begin()
+			tx.Find(&order)
+			if order.OrderNo != "9912" && order.OrderNo != "9913" {
+				t.Fatalf("idx: %v: order should comes from read db, but got order %v", j, order.OrderNo)
+			}
+			tx.Rollback()
+
+			tx = DB.Clauses(dbresolver.Write).Begin()
+			tx.Find(&order)
+			if order.OrderNo != "9911" {
+				t.Fatalf("idx: %v: order should comes from write db, but got order %v", j, order.OrderNo)
+			}
+			tx.Rollback()
+
+			tx = DB.Clauses(dbresolver.Use("users"), dbresolver.Write).Begin()
+			tx.Find(&order)
+			if order.OrderNo != "9914" {
+				t.Fatalf("idx: %v: order should comes from users, write db, but got order %v", j, order.OrderNo)
+			}
+			tx.Rollback()
+
+			tx = DB.Clauses(dbresolver.Write, dbresolver.Use("users")).Begin()
+			tx.Find(&order)
+			if order.OrderNo != "9914" {
+				t.Fatalf("idx: %v: order should comes from users, write db, but got order %v", j, order.OrderNo)
+			}
+			tx.Rollback()
+
+			// test query
 			DB.First(&order)
 			if order.OrderNo != "9912" && order.OrderNo != "9913" {
 				t.Fatalf("idx: %v: order should comes from read db, but got order %v", j, order.OrderNo)