From fef42941ba87bff8dad5d48b057a2c2056345984 Mon Sep 17 00:00:00 2001
From: qqxhb <30866940+qqxhb@users.noreply.github.com>
Date: Sat, 19 Aug 2023 21:33:31 +0800
Subject: [PATCH] feat: rm GetDBConnWithContext method (#6535)

* feat: rm contextconnpool method

* feat: nil
---
 go.sum          |  2 --
 gorm.go         |  9 ++++++---
 interfaces.go   |  6 ------
 prepare_stmt.go | 23 ++++++++++++++++++-----
 4 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/go.sum b/go.sum
index fb4240eb..bd6104c9 100644
--- a/go.sum
+++ b/go.sum
@@ -1,6 +1,4 @@
 github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
 github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
-github.com/jinzhu/now v1.1.4 h1:tHnRBy1i5F2Dh8BAFxqFzxKqqvezXrL2OW1TnX+Mlas=
-github.com/jinzhu/now v1.1.4/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
 github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
 github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
diff --git a/gorm.go b/gorm.go
index 203527af..32193870 100644
--- a/gorm.go
+++ b/gorm.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"reflect"
 	"sort"
 	"sync"
 	"time"
@@ -374,9 +375,11 @@ func (db *DB) AddError(err error) error {
 // DB returns `*sql.DB`
 func (db *DB) DB() (*sql.DB, error) {
 	connPool := db.ConnPool
-
-	if connector, ok := connPool.(GetDBConnectorWithContext); ok && connector != nil {
-		return connector.GetDBConnWithContext(db)
+	if db.Statement != nil && db.Statement.ConnPool != nil {
+		connPool = db.Statement.ConnPool
+	}
+	if tx, ok := connPool.(*sql.Tx); ok && tx != nil {
+		return (*sql.DB)(reflect.ValueOf(tx).Elem().FieldByName("db").UnsafePointer()), nil
 	}
 
 	if dbConnector, ok := connPool.(GetDBConnector); ok && dbConnector != nil {
diff --git a/interfaces.go b/interfaces.go
index 1950d740..3bcc3d57 100644
--- a/interfaces.go
+++ b/interfaces.go
@@ -77,12 +77,6 @@ type GetDBConnector interface {
 	GetDBConn() (*sql.DB, error)
 }
 
-// GetDBConnectorWithContext represents SQL db connector which takes into
-// account the current database context
-type GetDBConnectorWithContext interface {
-	GetDBConnWithContext(db *DB) (*sql.DB, error)
-}
-
 // Rows rows interface
 type Rows interface {
 	Columns() ([]string, error)
diff --git a/prepare_stmt.go b/prepare_stmt.go
index 9d98c86e..aa944624 100644
--- a/prepare_stmt.go
+++ b/prepare_stmt.go
@@ -30,15 +30,11 @@ func NewPreparedStmtDB(connPool ConnPool) *PreparedStmtDB {
 	}
 }
 
-func (db *PreparedStmtDB) GetDBConnWithContext(gormdb *DB) (*sql.DB, error) {
+func (db *PreparedStmtDB) GetDBConn() (*sql.DB, error) {
 	if sqldb, ok := db.ConnPool.(*sql.DB); ok {
 		return sqldb, nil
 	}
 
-	if connector, ok := db.ConnPool.(GetDBConnectorWithContext); ok && connector != nil {
-		return connector.GetDBConnWithContext(gormdb)
-	}
-
 	if dbConnector, ok := db.ConnPool.(GetDBConnector); ok && dbConnector != nil {
 		return dbConnector.GetDBConn()
 	}
@@ -131,6 +127,19 @@ func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (Conn
 		tx, err := beginner.BeginTx(ctx, opt)
 		return &PreparedStmtTX{PreparedStmtDB: db, Tx: tx}, err
 	}
+
+	beginner, ok := db.ConnPool.(ConnPoolBeginner)
+	if !ok {
+		return nil, ErrInvalidTransaction
+	}
+
+	connPool, err := beginner.BeginTx(ctx, opt)
+	if err != nil {
+		return nil, err
+	}
+	if tx, ok := connPool.(Tx); ok {
+		return &PreparedStmtTX{PreparedStmtDB: db, Tx: tx}, nil
+	}
 	return nil, ErrInvalidTransaction
 }
 
@@ -176,6 +185,10 @@ type PreparedStmtTX struct {
 	PreparedStmtDB *PreparedStmtDB
 }
 
+func (db *PreparedStmtTX) GetDBConn() (*sql.DB, error) {
+	return db.PreparedStmtDB.GetDBConn()
+}
+
 func (tx *PreparedStmtTX) Commit() error {
 	if tx.Tx != nil && !reflect.ValueOf(tx.Tx).IsNil() {
 		return tx.Tx.Commit()