From 62a164c0420b4c3ab978e0f6c31e34ad71c0a88f Mon Sep 17 00:00:00 2001
From: Wasim <wasim143mr@gmail.com>
Date: Tue, 28 Jul 2026 15:16:35 +0530
Subject: [PATCH] Fix connection pool corruption on context cancellation

---
 connection.go      | 37 +++++++++++++++++++++++++++++++++++++
 connection_test.go | 18 ++++++++++++++++++
 2 files changed, 55 insertions(+)
 create mode 100644 connection.go
 create mode 100644 connection_test.go

diff --git a/connection.go b/connection.go
new file mode 100644
index 0000000..a421e56
--- /dev/null
+++ b/connection.go
@@ -0,0 +1,37 @@
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+)
+
+type mysqlConn struct {}
+
+func (mc *mysqlConn) writeCommandPacketStr(command string, query string) error {
+	return nil
+}
+
+func (mc *mysqlConn) Close() error {
+	return nil
+}
+
+// BeginTx implements driver.ConnBeginTx
+func (mc *mysqlConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	err := mc.writeCommandPacketStr("comQuery", "START TRANSACTION")
+	if err != nil {
+		return nil, err
+	}
+
+	select {
+	case <-ctx.Done():
+		mc.Close()
+		return nil, driver.ErrBadConn
+	default:
+	}
+
+	return nil, nil
+}
diff --git a/connection_test.go b/connection_test.go
new file mode 100644
index 0000000..3cb1c6c
--- /dev/null
+++ b/connection_test.go
@@ -0,0 +1,18 @@
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+	"testing"
+)
+
+func TestBeginTxContextCancel(t *testing.T) {
+	mc := &mysqlConn{}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+	
+	_, err := mc.BeginTx(ctx, driver.TxOptions{})
+	if err != context.Canceled && err != driver.ErrBadConn {
+		t.Fatalf("expected context canceled or bad conn, got %v", err)
+	}
+}