From 661781a3d7a36b33d8a30d97499d7b428939a9ec Mon Sep 17 00:00:00 2001
From: Lev Zakharov <l.j.zakharov@gmail.com>
Date: Mon, 5 Jun 2023 11:25:05 +0300
Subject: [PATCH] feat: add *sql.DB connector that uses database context
 (#6366)

* feat: add SQLConnector

* rename
---
 gorm.go       | 4 ++++
 interfaces.go | 6 ++++++
 2 files changed, 10 insertions(+)

diff --git a/gorm.go b/gorm.go
index 21b289db..9297850e 100644
--- 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
index 3bcc3d57..1950d740 100644
--- 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)