From 201d743eee94c1537869f230cb31f9ed64b55949 Mon Sep 17 00:00:00 2001 From: Owen Smallwood Date: Fri, 12 Dec 2025 14:21:01 -0600 Subject: [PATCH] adds tests --- pkg/storage/unified/sql/db/dbimpl/db_test.go | 165 +++++++++++++++++++ 1 file changed, 165 insertions(+) diff --git a/pkg/storage/unified/sql/db/dbimpl/db_test.go b/pkg/storage/unified/sql/db/dbimpl/db_test.go index 6006e79d7f4..059de9c0255 100644 --- a/pkg/storage/unified/sql/db/dbimpl/db_test.go +++ b/pkg/storage/unified/sql/db/dbimpl/db_test.go @@ -1,8 +1,13 @@ package dbimpl import ( + "context" "database/sql" + "database/sql/driver" + "sync" + "sync/atomic" "testing" + "time" "github.com/stretchr/testify/require" @@ -25,3 +30,163 @@ func TestDB_BeginTx(t *testing.T) { require.NoError(t, err) require.NotNil(t, tx) } + +func TestDB_BeginTx_RetriesOnTimeout(t *testing.T) { + t.Parallel() + + // Register a driver that times out on the first two connection attempts + driverName := "test-timeout-driver-" + t.Name() + timeoutDriver := &timeoutTestDriver{ + timeoutsRemaining: 2, + } + sql.Register(driverName, timeoutDriver) + + sqlDB, err := sql.Open(driverName, "") + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + + d := NewDB(sqlDB, driverName) + + // Use a longer timeout context since we expect retries with backoff + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + tx, err := d.BeginTx(ctx, nil) + require.NoError(t, err) + require.NotNil(t, tx) + + // Should have attempted 3 connections (2 timeouts + 1 success) + require.Equal(t, int32(3), timeoutDriver.connAttempts.Load()) +} + +func TestDB_BeginTx_ConnectionReleasedAfterCommitAndRollback(t *testing.T) { + t.Parallel() + + // Register a driver that tracks open connections + driverName := "test-conn-tracking-driver-" + t.Name() + trackingDriver := &connTrackingDriver{} + sql.Register(driverName, trackingDriver) + + sqlDB, err := sql.Open(driverName, "") + require.NoError(t, err) + // Setting MaxIdleConns to 0 ensures connections are closed when returned to pool + sqlDB.SetMaxIdleConns(0) + sqlDB.SetMaxOpenConns(1) + + d := NewDB(sqlDB, driverName) + ctx := testutil.NewDefaultTestContext(t) + + // Test Commit releases connection + tx1, err := d.BeginTx(ctx, nil) + require.NoError(t, err) + require.Equal(t, int32(1), trackingDriver.openConns.Load(), "should have 1 open connection after BeginTx") + + err = tx1.Commit() + require.NoError(t, err) + require.Equal(t, int32(0), trackingDriver.openConns.Load(), "should have 0 open connections after Commit") + + // Test Rollback releases connection + tx2, err := d.BeginTx(ctx, nil) + require.NoError(t, err) + require.Equal(t, int32(1), trackingDriver.openConns.Load(), "should have 1 open connection after BeginTx") + + err = tx2.Rollback() + require.NoError(t, err) + require.Equal(t, int32(0), trackingDriver.openConns.Load(), "should have 0 open connections after Rollback") +} + +// timeoutTestDriver simulates connection timeouts for testing retry logic +type timeoutTestDriver struct { + mu sync.Mutex + timeoutsRemaining int + connAttempts atomic.Int32 +} + +func (d *timeoutTestDriver) Open(name string) (driver.Conn, error) { + return d.connect(context.Background()) +} + +func (d *timeoutTestDriver) OpenConnector(name string) (driver.Connector, error) { + return &timeoutTestConnector{driver: d}, nil +} + +func (d *timeoutTestDriver) connect(ctx context.Context) (driver.Conn, error) { + d.connAttempts.Add(1) + + d.mu.Lock() + shouldTimeout := d.timeoutsRemaining > 0 + if shouldTimeout { + d.timeoutsRemaining-- + } + d.mu.Unlock() + + if shouldTimeout { + // Block until context is cancelled to simulate a slow connection + <-ctx.Done() + return nil, ctx.Err() + } + + return &timeoutTestConn{}, nil +} + +type timeoutTestConnector struct { + driver *timeoutTestDriver +} + +func (c *timeoutTestConnector) Connect(ctx context.Context) (driver.Conn, error) { + return c.driver.connect(ctx) +} + +func (c *timeoutTestConnector) Driver() driver.Driver { + return c.driver +} + +type timeoutTestConn struct{} + +func (c *timeoutTestConn) Prepare(query string) (driver.Stmt, error) { return testStmt{}, nil } +func (c *timeoutTestConn) Close() error { return nil } +func (c *timeoutTestConn) Begin() (driver.Tx, error) { return testTx{}, nil } +func (c *timeoutTestConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + return testTx{}, nil +} + +// connTrackingDriver tracks open connections for testing connection lifecycle +type connTrackingDriver struct { + openConns atomic.Int32 +} + +func (d *connTrackingDriver) Open(name string) (driver.Conn, error) { + d.openConns.Add(1) + return &connTrackingConn{driver: d}, nil +} + +func (d *connTrackingDriver) OpenConnector(name string) (driver.Connector, error) { + return &connTrackingConnector{driver: d}, nil +} + +type connTrackingConnector struct { + driver *connTrackingDriver +} + +func (c *connTrackingConnector) Connect(ctx context.Context) (driver.Conn, error) { + c.driver.openConns.Add(1) + return &connTrackingConn{driver: c.driver}, nil +} + +func (c *connTrackingConnector) Driver() driver.Driver { + return c.driver +} + +type connTrackingConn struct { + driver *connTrackingDriver +} + +func (c *connTrackingConn) Prepare(query string) (driver.Stmt, error) { return testStmt{}, nil } +func (c *connTrackingConn) Close() error { + c.driver.openConns.Add(-1) + return nil +} +func (c *connTrackingConn) Begin() (driver.Tx, error) { return testTx{}, nil } +func (c *connTrackingConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + return testTx{}, nil +}