Add context-based transaction support with EnsureTx

Add EnsureTx, WithTransaction, and WithoutTransaction for managing
transactions via context. getTxOrDB now checks ctx for transactions,
so WithContext(ctx) automatically uses the transaction if present.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

database.go

@@ -26,6 +26,10 @@ type Database interface {
 	GetDB() *sql.DB
 	// BeginTx starts a transaction and executes the function f.
 	BeginTx(ctx context.Context, opts *sql.TxOptions, f func(tx Transaction) error) error
+	// EnsureTx ensures the function f runs within a transaction.
+	// If ctx already contains a transaction started by a previous EnsureTx call, it reuses that transaction.
+	// Otherwise, it begins a new transaction and stores it in the context.
+	EnsureTx(ctx context.Context, opts *sql.TxOptions, f func(ctx context.Context) error) error
 	// Query executes a query and returns the cursor
 	Query(sql string) (Cursor, error)
 	// QueryContext executes a query with context and returns the cursor
@@ -187,10 +191,15 @@ func (d database) GetDB() *sql.DB {
 	return d.db
 }
 
-func (d database) getTxOrDB() txOrDB {
+func (d database) getTxOrDB(ctx context.Context) txOrDB {
 	if d.tx != nil {
 		return d.tx
 	}
+	if ctx != nil {
+		if tx, ok := ctx.Value(txContextKey{}).(Transaction); ok {
+			return tx.GetTx()
+		}
+	}
 	return d.db
 }
 
@@ -229,7 +238,7 @@ func (d database) queryContextOnce(ctx context.Context, sqlString string, retry
 	interceptor := d.interceptor
 	var rows *sql.Rows
 	invoker := func(ctx context.Context, sql string) (err error) {
-		rows, err = d.getTxOrDB().QueryContext(ctx, sql)
+		rows, err = d.getTxOrDB(ctx).QueryContext(ctx, sql)
 		return
 	}
 
@@ -266,7 +275,7 @@ func (d database) ExecuteContext(ctx context.Context, sqlString string) (sql.Res
 
 	var result sql.Result
 	invoker := func(ctx context.Context, sql string) (err error) {
-		result, err = d.getTxOrDB().ExecContext(ctx, sql)
+		result, err = d.getTxOrDB(ctx).ExecContext(ctx, sql)
 		return
 	}
 	var err error

transaction.go

@@ -21,6 +21,18 @@ type Transaction interface {
 	DeleteFrom(table Table) deleteWithTable
 }
 
+type txContextKey struct{}
+
+// WithTransaction stores the transaction in the context.
+func WithTransaction(ctx context.Context, tx Transaction) context.Context {
+	return context.WithValue(ctx, txContextKey{}, tx)
+}
+
+// WithoutTransaction returns a context without the transaction.
+func WithoutTransaction(ctx context.Context) context.Context {
+	return context.WithValue(ctx, txContextKey{}, nil)
+}
+
 func (d *database) GetTx() *sql.Tx {
 	return d.tx
 }
@@ -56,3 +68,21 @@ func (d *database) BeginTx(ctx context.Context, opts *sql.TxOptions, f func(tx T
 	isCommitted = true
 	return nil
 }
+
+// EnsureTx ensures the function f runs within a transaction.
+// If ctx already contains a transaction started by a previous EnsureTx call, it reuses that transaction.
+// Otherwise, it begins a new transaction and stores it in the context.
+func (d *database) EnsureTx(ctx context.Context, opts *sql.TxOptions, f func(ctx context.Context) error) error {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	if _, ok := ctx.Value(txContextKey{}).(Transaction); ok {
+		return f(ctx)
+	}
+	if d.tx != nil {
+		return f(WithTransaction(ctx, d))
+	}
+	return d.BeginTx(ctx, opts, func(tx Transaction) error {
+		return f(WithTransaction(ctx, tx))
+	})
+}

transaction_test.go

@@ -81,3 +81,152 @@ func TestTransaction(t *testing.T) {
 		t.Error("should get error here")
 	}
 }
+
+func TestWithTransaction(t *testing.T) {
+	ctx := context.Background()
+
+	// ctx without transaction
+	if _, ok := ctx.Value(txContextKey{}).(Transaction); ok {
+		t.Error("should not have transaction")
+	}
+
+	db := newMockDatabase()
+	err := db.BeginTx(ctx, nil, func(tx Transaction) error {
+		txCtx := WithTransaction(ctx, tx)
+		got, ok := txCtx.Value(txContextKey{}).(Transaction)
+		if !ok {
+			t.Error("should have transaction")
+		}
+		if got.GetTx() != tx.GetTx() {
+			t.Error("should be the same tx")
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestWithoutTransaction(t *testing.T) {
+	db := newMockDatabase()
+	err := db.BeginTx(context.Background(), nil, func(tx Transaction) error {
+		txCtx := WithTransaction(context.Background(), tx)
+		cleanCtx := WithoutTransaction(txCtx)
+		if _, ok := cleanCtx.Value(txContextKey{}).(Transaction); ok {
+			t.Error("should not have transaction after WithoutTransaction")
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestEnsureTx(t *testing.T) {
+	db := newMockDatabase()
+	sharedMockConn.mockTx = nil
+
+	// EnsureTx should create a new transaction
+	err := db.EnsureTx(context.Background(), nil, func(ctx context.Context) error {
+		tx, ok := ctx.Value(txContextKey{}).(Transaction)
+		if !ok || tx.GetTx() == nil {
+			t.Error("should have transaction in ctx")
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+	if !sharedMockConn.mockTx.isCommitted {
+		t.Error("should be committed")
+	}
+
+	// EnsureTx with nil ctx
+	err = db.EnsureTx(nil, nil, func(ctx context.Context) error {
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+
+	// EnsureTx should rollback on error
+	err = db.EnsureTx(context.Background(), nil, func(ctx context.Context) error {
+		return errors.New("error")
+	})
+	if err == nil {
+		t.Error("should get error here")
+	}
+	if !sharedMockConn.mockTx.isRolledBack {
+		t.Error("should be rolled back")
+	}
+
+	// EnsureTx should reuse transaction from ctx
+	sharedMockConn.mockTx = nil
+	err = db.EnsureTx(context.Background(), nil, func(ctx context.Context) error {
+		outerTx := ctx.Value(txContextKey{}).(Transaction)
+
+		// nested EnsureTx should reuse the same tx
+		return db.EnsureTx(ctx, nil, func(innerCtx context.Context) error {
+			innerTx := innerCtx.Value(txContextKey{}).(Transaction)
+			if innerTx.GetTx() != outerTx.GetTx() {
+				t.Error("nested EnsureTx should reuse the same transaction")
+			}
+			return nil
+		})
+	})
+	if err != nil {
+		t.Error(err)
+	}
+
+	// EnsureTx should reuse transaction from BeginTx
+	err = db.BeginTx(context.Background(), nil, func(tx Transaction) error {
+		txDb := tx.(*database)
+		return txDb.EnsureTx(context.Background(), nil, func(ctx context.Context) error {
+			ctxTx, ok := ctx.Value(txContextKey{}).(Transaction)
+			if !ok {
+				t.Error("should have transaction in ctx")
+			}
+			if ctxTx.GetTx() != tx.GetTx() {
+				t.Error("EnsureTx inside BeginTx should reuse the same transaction")
+			}
+			return nil
+		})
+	})
+	if err != nil {
+		t.Error(err)
+	}
+
+	// EnsureTx should fail if BeginTx fails
+	sharedMockConn.beginTxError = errors.New("error")
+	err = db.EnsureTx(context.Background(), nil, func(ctx context.Context) error {
+		return nil
+	})
+	if err == nil {
+		t.Error("should get error here")
+	}
+	sharedMockConn.beginTxError = nil
+}
+
+func TestGetTxOrDBWithContext(t *testing.T) {
+	db := newMockDatabase()
+
+	// without tx in ctx, should return db
+	d := db.(*database)
+	result := d.getTxOrDB(context.Background())
+	if result != d.db {
+		t.Error("should return db when no tx in context")
+	}
+
+	// with tx in ctx, should return tx
+	err := db.BeginTx(context.Background(), nil, func(tx Transaction) error {
+		txCtx := WithTransaction(context.Background(), tx)
+		result := d.getTxOrDB(txCtx)
+		if result != tx.GetTx() {
+			t.Error("should return tx from context")
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}