feat: add *sql.DB connector that uses database context (#6366)

* feat: add SQLConnector * rename
go-gorm · Jun 5, 2023 · 661781a · 661781a
1 parent 5eaccaa
commit 661781a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/gorm.go b/gorm.go
@@ -375,6 +375,10 @@ func (db *DB) AddError(err error) error {
 func (db *DB) DB() (*sql.DB, error) {
 	connPool := db.ConnPool
 
+	if connector, ok := connPool.(GetDBConnectorWithContext); ok && connector != nil {
+		return connector.GetDBConnWithContext(db)
+	}
+
 	if dbConnector, ok := connPool.(GetDBConnector); ok && dbConnector != nil {
 		if sqldb, err := dbConnector.GetDBConn(); sqldb != nil || err != nil {
 			return sqldb, err

diff --git a/interfaces.go b/interfaces.go
@@ -77,6 +77,12 @@ 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)