diff --git a/finisher_api.go b/finisher_api.go index d38d60b7..dd0eb83a 100644 --- a/finisher_api.go +++ b/finisher_api.go @@ -515,6 +515,30 @@ func (db *DB) ScanRows(rows *sql.Rows, dest interface{}) error { return tx.Error } +// Connection use a db conn to execute Multiple commands,this conn will put conn pool after it is executed. +func (db *DB) Connection(fc func(tx *DB) error) (err error) { + if db.Error != nil { + return db.Error + } + + tx := db.getInstance() + sqlDB, err := tx.DB() + if err != nil { + return + } + + conn, err := sqlDB.Conn(tx.Statement.Context) + if err != nil { + return + } + + defer conn.Close() + tx.Statement.ConnPool = conn + err = fc(tx) + + return +} + // Transaction start a transaction as a block, return error will rollback, otherwise to commit. func (db *DB) Transaction(fc func(tx *DB) error, opts ...*sql.TxOptions) (err error) { panicked := true diff --git a/tests/connection_test.go b/tests/connection_test.go new file mode 100644 index 00000000..9b5dcd05 --- /dev/null +++ b/tests/connection_test.go @@ -0,0 +1,48 @@ +package tests_test + +import ( + "fmt" + "gorm.io/driver/mysql" + "gorm.io/gorm" + "testing" +) + +func TestWithSingleConnection(t *testing.T) { + + var expectedName = "test" + var actualName string + + setSQL, getSQL := getSetSQL(DB.Dialector.Name()) + if len(setSQL) == 0 || len(getSQL) == 0 { + return + } + + err := DB.Connection(func(tx *gorm.DB) error { + if err := tx.Exec(setSQL, expectedName).Error; err != nil { + return err + } + + if err := tx.Raw(getSQL).Scan(&actualName).Error; err != nil { + return err + } + return nil + }) + + if err != nil { + t.Errorf(fmt.Sprintf("WithSingleConnection should work, but got err %v", err)) + } + + if actualName != expectedName { + t.Errorf("WithSingleConnection() method should get correct value, expect: %v, got %v", expectedName, actualName) + } + +} + +func getSetSQL(driverName string) (string, string) { + switch driverName { + case mysql.Dialector{}.Name(): + return "SET @testName := ?", "SELECT @testName" + default: + return "", "" + } +}