diff --git a/connection.go b/connection.go new file mode 100644 index 0000000..ef8e2ad --- /dev/null +++ b/connection.go @@ -0,0 +1,177 @@ +// Go MySQL Driver - A MySQL-Driver for Go's database/sql package +// +// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this file, +// You can obtain one at http://mozilla.org/MPL/2.0/. + +package mysql + +import ( + "context" + "database/sql/driver" + "net" + "time" +) + +type mysqlConn struct { + netConn net.Conn + affectedRows uint64 + insertId uint64 + cfg *Config + maxAllowedPacket int + maxWriteSize int + writeTimeout time.Duration + flags clientFlag + status statusFlag + sequence uint8 + parseTime bool + strict bool + buf buffer + closed bool +} + +type mysqlTx struct { + mc *mysqlConn +} + +func (mc *mysqlConn) Begin() (driver.Tx, error) { + return mc.begin(false) +} + +func (mc *mysqlConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) { + if mc.closed { + return nil, driver.ErrBadConn + } + + // Check if context is already canceled before starting + if err := ctx.Err(); err != nil { + return nil, err + } + + var query string + if opts.ReadOnly { + query = "START TRANSACTION READ ONLY" + } else { + switch sql.IsolationLevel(opts.Isolation) { + case sql.LevelDefault: + query = "START TRANSACTION" + case sql.LevelReadUncommitted: + query = "SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED; START TRANSACTION" + case sql.LevelReadCommitted: + query = "SET TRANSACTION ISOLATION LEVEL READ COMMITTED; START TRANSACTION" + case sql.LevelRepeatableRead: + query = "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ; START TRANSACTION" + case sql.LevelSerializable: + query = "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE; START TRANSACTION" + default: + return nil, driver.ErrBadConn + } + } + + // Send the START TRANSACTION command + err := mc.exec(query) + if err != nil { + return nil, err + } + + // Critical: Check if context was canceled after sending the command + // If canceled, the transaction may have started on the server side + // We must close the connection to prevent returning a dirty connection to the pool + select { + case <-ctx.Done(): + // Transaction started on server but context is dead + // Mark connection as bad and close it + mc.cleanup() + return nil, ctx.Err() + default: + // Context still valid, proceed normally + } + + return &mysqlTx{mc}, nil +} + +func (mc *mysqlConn) begin(readOnly bool) (driver.Tx, error) { + if mc.closed { + return nil, driver.ErrBadConn + } + var query string + if readOnly { + query = "START TRANSACTION READ ONLY" + } else { + query = "START TRANSACTION" + } + err := mc.exec(query) + if err != nil { + return nil, err + } + return &mysqlTx{mc}, nil +} + +func (mc *mysqlConn) Close() error { + if mc.closed { + return nil + } + mc.cleanup() + return nil +} + +func (mc *mysqlConn) cleanup() { + if !mc.closed { + mc.closed = true + if mc.netConn != nil { + mc.netConn.Close() + } + } +} + +func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) { + if mc.closed { + return nil, driver.ErrBadConn + } + return mc.prepare(query) +} + +func (mc *mysqlConn) prepare(query string) (driver.Stmt, error) { + // Implementation placeholder + return nil, nil +} + +func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, error) { + if mc.closed { + return nil, driver.ErrBadConn + } + return mc.exec(query) +} + +func (mc *mysqlConn) exec(query string) error { + // Implementation placeholder + return nil +} + +func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, error) { + if mc.closed { + return nil, driver.ErrBadConn + } + return mc.query(query) +} + +func (mc *mysqlConn) query(query string) (driver.Rows, error) { + // Implementation placeholder + return nil, nil +} + +func (tx *mysqlTx) Commit() error { + if tx.mc.closed { + return driver.ErrBadConn + } + return tx.mc.exec("COMMIT") +} + +func (tx *mysqlTx) Rollback() error { + if tx.mc.closed { + return driver.ErrBadConn + } + return tx.mc.exec("ROLLBACK") +} diff --git a/connection_test.go b/connection_test.go new file mode 100644 index 0000000..0926e6f --- /dev/null +++ b/connection_test.go @@ -0,0 +1,155 @@ +// Go MySQL Driver - A MySQL-Driver for Go's database/sql package +// +// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this file, +// You can obtain one at http://mozilla.org/MPL/2.0/. + +package mysql + +import ( + "context" + "database/sql" + "database/sql/driver" + "testing" + "time" +) + +// TestBeginTx_ContextCanceledBeforeStart verifies that a pre-canceled context +// causes BeginTx to fail without leaking a connection +func TestBeginTx_ContextCanceledBeforeStart(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + mc := &mysqlConn{ + closed: false, + } + + tx, err := mc.BeginTx(ctx, driver.TxOptions{}) + if err == nil { + tx.Rollback() + t.Fatal("Expected BeginTx to fail with canceled context") + } + + if err != context.Canceled { + t.Errorf("Expected context.Canceled error, got: %v", err) + } +} + +// TestBeginTx_ContextCanceledDuringExecution simulates a context cancellation +// that occurs during the transaction start, verifying the connection is closed +func TestBeginTx_ContextCanceledDuringExecution(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond) + defer cancel() + + // Give the context time to expire + time.Sleep(10 * time.Millisecond) + + mc := &mysqlConn{ + closed: false, + } + + tx, err := mc.BeginTx(ctx, driver.TxOptions{}) + if err == nil { + tx.Rollback() + t.Fatal("Expected BeginTx to fail with deadline exceeded") + } + + if ctx.Err() == nil { + t.Fatal("Context should be canceled") + } + + // Verify connection was closed to prevent pool contamination + if !mc.closed { + t.Error("Connection should be closed after context cancellation during BeginTx") + } +} + +// TestBeginTx_SuccessfulTransaction verifies normal transaction flow +func TestBeginTx_SuccessfulTransaction(t *testing.T) { + ctx := context.Background() + + mc := &mysqlConn{ + closed: false, + } + + tx, err := mc.BeginTx(ctx, driver.TxOptions{}) + if err != nil { + t.Fatalf("BeginTx should succeed with valid context: %v", err) + } + + if tx == nil { + t.Fatal("Transaction should not be nil") + } + + // Verify connection is still open + if mc.closed { + t.Error("Connection should remain open after successful BeginTx") + } + + err = tx.Commit() + if err != nil { + t.Errorf("Commit failed: %v", err) + } +} + +// TestBeginTx_ReadOnlyTransaction verifies read-only transaction option +func TestBeginTx_ReadOnlyTransaction(t *testing.T) { + ctx := context.Background() + + mc := &mysqlConn{ + closed: false, + } + + tx, err := mc.BeginTx(ctx, driver.TxOptions{ReadOnly: true}) + if err != nil { + t.Fatalf("BeginTx with ReadOnly should succeed: %v", err) + } + + if tx == nil { + t.Fatal("Transaction should not be nil") + } + + err = tx.Rollback() + if err != nil { + t.Errorf("Rollback failed: %v", err) + } +} + +// TestBeginTx_IsolationLevels verifies different isolation levels +func TestBeginTx_IsolationLevels(t *testing.T) { + testCases := []struct { + name string + isolation driver.IsolationLevel + }{ + {"Default", driver.IsolationLevel(sql.LevelDefault)}, + {"ReadUncommitted", driver.IsolationLevel(sql.LevelReadUncommitted)}, + {"ReadCommitted", driver.IsolationLevel(sql.LevelReadCommitted)}, + {"RepeatableRead", driver.IsolationLevel(sql.LevelRepeatableRead)}, + {"Serializable", driver.IsolationLevel(sql.LevelSerializable)}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + mc := &mysqlConn{ + closed: false, + } + + tx, err := mc.BeginTx(ctx, driver.TxOptions{Isolation: tc.isolation}) + if err != nil { + t.Fatalf("BeginTx with %s isolation should succeed: %v", tc.name, err) + } + + if tx == nil { + t.Fatal("Transaction should not be nil") + } + + err = tx.Rollback() + if err != nil { + t.Errorf("Rollback failed: %v", err) + } + }) + } +}