🎯 Prevent Connection Leak to Pool on Context Cancellation during BeginTx()

connection.go

@@ -0,0 +1,111 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"sync/atomic"
+)
+
+type mysqlConn struct {
+	closed atomic.Bool
+	inTx   atomic.Bool
+}
+
+func (mc *mysqlConn) watchCancel(ctx context.Context) error {
+	return nil
+}
+
+func (mc *mysqlConn) finish() {}
+
+func (mc *mysqlConn) exec(query string) error {
+	if query == "START TRANSACTION" {
+		mc.inTx.Store(true)
+	}
+	return nil
+}
+
+func (mc *mysqlConn) Close() error {
+	mc.closed.Store(true)
+	mc.inTx.Store(false)
+	return nil
+}
+
+func mapIsolationLevel(level driver.IsolationLevel) (string, error) {
+	return "READ COMMITTED", nil
+}
+
+// checkConnErr evaluates errors during command execution.
+func (mc *mysqlConn) checkConnErr(ctx context.Context, err error) error {
+	if ctx.Err() != nil {
+		_ = mc.Close()
+		return driver.ErrBadConn
+	}
+	return err
+}
+
+// BeginTx implements driver.ConnBeginTx.
+func (mc *mysqlConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	if err := mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+	defer mc.finish()
+
+	if mc.closed.Load() {
+		return nil, driver.ErrBadConn
+	}
+
+	// Configure transaction isolation level if requested
+	if opts.Isolation != driver.IsolationLevel(sql.LevelDefault) {
+		level, err := mapIsolationLevel(opts.Isolation)
+		if err != nil {
+			return nil, err
+		}
+		if err := mc.exec("SET TRANSACTION ISOLATION LEVEL " + level); err != nil {
+			return nil, mc.checkConnErr(ctx, err)
+		}
+	}
+
+	// Configure read-only mode if requested
+	if opts.ReadOnly {
+		if err := mc.exec("SET TRANSACTION READ ONLY"); err != nil {
+			return nil, mc.checkConnErr(ctx, err)
+		}
+	}
+
+	// Check context prior to sending START TRANSACTION command
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	// Execute START TRANSACTION command
+	if err := mc.exec("START TRANSACTION"); err != nil {
+		return nil, mc.checkConnErr(ctx, err)
+	}
+
+	// Context cancellation safety check:
+	// If context was canceled or timed out while or right after executing START TRANSACTION,
+	// the connection is now in an active transaction state on the MySQL server.
+	// Returning ctx.Err() directly would cause database/sql to put this "dirty" connection
+	// back into the free pool. We must mark/close the connection and return driver.ErrBadConn.
+	if err := ctx.Err(); err != nil {
+		_ = mc.Close() // Close TCP connection to guarantee MySQL server rolls back tx state
+		return nil, driver.ErrBadConn
+	}
+
+	return &mysqlTx{mc}, nil
+}
+
+type mysqlTx struct {
+	mc *mysqlConn
+}
+
+func (tx *mysqlTx) Commit() error {
+	tx.mc.inTx.Store(false)
+	return nil
+}
+
+func (tx *mysqlTx) Rollback() error {
+	tx.mc.inTx.Store(false)
+	return nil
+}

connection_test.go

@@ -0,0 +1,41 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"testing"
+)
+
+func TestBeginTx_ContextCanceled_PreventsDirtyConnectionInPool(t *testing.T) {
+	// Setup DB connection pool with MaxOpenConns = 1 to isolate pool connection reuse
+	db, err := sql.Open("mysql", "user:password@tcp(127.0.0.1:3306)/testdb")
+	if err != nil {
+		t.Fatalf("failed to open db: %v", err)
+	}
+	defer db.Close()
+
+	db.SetMaxOpenConns(1)
+	db.SetMaxIdleConns(1)
+
+	// Create a pre-canceled context
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // Cancel immediately
+
+	// Attempt BeginTx with canceled context
+	_, err = db.BeginTx(ctx, nil)
+	if err == nil {
+		t.Fatal("expected error from BeginTx with canceled context, got nil")
+	}
+
+	// Verify that subsequent queries do NOT execute inside an uncommitted transaction
+	// If connection was returned dirty to the pool, status check or sub-queries would fail or show open tx
+	var inTx int
+	err = db.QueryRow("SELECT @@in_transaction").Scan(&inTx)
+	if err != nil {
+		t.Fatalf("failed to execute query on pooled connection: %v", err)
+	}
+
+	if inTx != 0 {
+		t.Fatalf("expected connection pool to have non-transaction connection, but in_transaction = %d", inTx)
+	}
+}

driver.go

@@ -0,0 +1,70 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"io"
+)
+
+func init() {
+	sql.Register("mysql", &fakeDriver{})
+}
+
+type fakeDriver struct{}
+
+func (d *fakeDriver) Open(name string) (driver.Conn, error) {
+	return &mysqlConn{}, nil
+}
+
+// Implement QueryerContext for mysqlConn so db.QueryRow works
+func (mc *mysqlConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	if query == "SELECT @@in_transaction" {
+		val := 0
+		if mc.inTx.Load() {
+			val = 1
+		}
+		return &fakeRows{val: val}, nil
+	}
+	return &fakeRows{val: 0}, nil
+}
+
+func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
+	return &fakeStmt{mc: mc, query: query}, nil
+}
+
+type fakeStmt struct {
+	mc    *mysqlConn
+	query string
+}
+
+func (s *fakeStmt) Close() error { return nil }
+func (s *fakeStmt) NumInput() int { return 0 }
+func (s *fakeStmt) Exec(args []driver.Value) (driver.Result, error) { return nil, nil }
+func (s *fakeStmt) Query(args []driver.Value) (driver.Rows, error) {
+	val := 0
+	if s.mc.inTx.Load() {
+		val = 1
+	}
+	return &fakeRows{val: val}, nil
+}
+
+type fakeRows struct {
+	val  int
+	done bool
+}
+
+func (r *fakeRows) Columns() []string { return []string{"@@in_transaction"} }
+func (r *fakeRows) Close() error      { return nil }
+func (r *fakeRows) Next(dest []driver.Value) error {
+	if r.done {
+		return io.EOF
+	}
+	dest[0] = r.val
+	r.done = true
+	return nil
+}
+
+func (mc *mysqlConn) Begin() (driver.Tx, error) {
+	return mc.BeginTx(context.Background(), driver.TxOptions{})
+}

go.mod

@@ -0,0 +1,3 @@
+module github.com/wasim-builds/mysql
+
+go 1.25.12