From a983af405611725baf80a3ffa85e4864d363917e Mon Sep 17 00:00:00 2001
From: qqxhb <30866940+qqxhb@users.noreply.github.com>
Date: Thu, 19 Jan 2023 10:07:17 +0800
Subject: [PATCH] fix: clone and write db (#748)

---
 do.go                               | 2 +-
 internal/template/query.go          | 4 ++--
 tests/.expect/dal_1/query/gen.go    | 4 ++--
 tests/.expect/dal_2/query/gen.go    | 4 ++--
 tests/.expect/dal_3/query/gen.go    | 4 ++--
 tests/.expect/dal_test/query/gen.go | 4 ++--
 6 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/do.go b/do.go
index adfdd656..f84b0e84 100644
--- a/do.go
+++ b/do.go
@@ -70,7 +70,7 @@ func (d *DO) ReplaceDB(db *gorm.DB) {
 
 // ReplaceConnPool replace db connection pool
 func (d *DO) ReplaceConnPool(pool gorm.ConnPool) {
-	d.db = d.db.Session(&gorm.Session{})
+	d.db = d.db.Session(&gorm.Session{Initialized: true}).Session(&gorm.Session{})
 	d.db.Statement.ConnPool = pool
 }
 
diff --git a/internal/template/query.go b/internal/template/query.go
index e67fa5de..cdd4d7c5 100644
--- a/internal/template/query.go
+++ b/internal/template/query.go
@@ -49,11 +49,11 @@ func (q *Query) clone(db *gorm.DB) *Query {
 }
 
 func (q *Query) ReadDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Read))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Read))
 }
 
 func (q *Query) WriteDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Write))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Write))
 }
 
 func (q *Query) ReplaceDB(db *gorm.DB) *Query {
diff --git a/tests/.expect/dal_1/query/gen.go b/tests/.expect/dal_1/query/gen.go
index 79573df5..cfa3e63e 100644
--- a/tests/.expect/dal_1/query/gen.go
+++ b/tests/.expect/dal_1/query/gen.go
@@ -68,11 +68,11 @@ func (q *Query) clone(db *gorm.DB) *Query {
 }
 
 func (q *Query) ReadDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Read))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Read))
 }
 
 func (q *Query) WriteDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Write))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Write))
 }
 
 func (q *Query) ReplaceDB(db *gorm.DB) *Query {
diff --git a/tests/.expect/dal_2/query/gen.go b/tests/.expect/dal_2/query/gen.go
index 79573df5..cfa3e63e 100644
--- a/tests/.expect/dal_2/query/gen.go
+++ b/tests/.expect/dal_2/query/gen.go
@@ -68,11 +68,11 @@ func (q *Query) clone(db *gorm.DB) *Query {
 }
 
 func (q *Query) ReadDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Read))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Read))
 }
 
 func (q *Query) WriteDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Write))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Write))
 }
 
 func (q *Query) ReplaceDB(db *gorm.DB) *Query {
diff --git a/tests/.expect/dal_3/query/gen.go b/tests/.expect/dal_3/query/gen.go
index 19fb643a..53b8245e 100644
--- a/tests/.expect/dal_3/query/gen.go
+++ b/tests/.expect/dal_3/query/gen.go
@@ -68,11 +68,11 @@ func (q *Query) clone(db *gorm.DB) *Query {
 }
 
 func (q *Query) ReadDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Read))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Read))
 }
 
 func (q *Query) WriteDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Write))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Write))
 }
 
 func (q *Query) ReplaceDB(db *gorm.DB) *Query {
diff --git a/tests/.expect/dal_test/query/gen.go b/tests/.expect/dal_test/query/gen.go
index 86fa5ede..845f24bb 100644
--- a/tests/.expect/dal_test/query/gen.go
+++ b/tests/.expect/dal_test/query/gen.go
@@ -66,11 +66,11 @@ func (q *Query) clone(db *gorm.DB) *Query {
 }
 
 func (q *Query) ReadDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Read))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Read))
 }
 
 func (q *Query) WriteDB() *Query {
-	return q.clone(q.db.Clauses(dbresolver.Write))
+	return q.ReplaceDB(q.db.Clauses(dbresolver.Write))
 }
 
 func (q *Query) ReplaceDB(db *gorm.DB) *Query {