diff --git a/connection.go b/connection.go new file mode 100644 index 0000000..6ce9c3e --- /dev/null +++ b/connection.go @@ -0,0 +1,109 @@ +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) + } + + // CRITICAL FIX: Check if context was canceled during the roundtrip + select { + case <-ctx.Done(): + mc.Close() + return nil, ctx.Err() + default: + } + + 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 +} diff --git a/connection_test.go b/connection_test.go new file mode 100644 index 0000000..5761837 --- /dev/null +++ b/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) + } +} diff --git a/driver.go b/driver.go new file mode 100644 index 0000000..a917908 --- /dev/null +++ b/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{}) +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..e212dc3 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module github.com/wasim-builds/mysql + +go 1.25.12