driver_test.go

package main

import (
	"context"
	"database/sql"
	"errors"
	"os"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

func TestTxCommitContextCancellation(t *testing.T) {
	connString := os.Getenv("DATABASE_URL")
	if connString == "" {
		connString = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
	}

	config, err := pgxpool.ParseConfig(connString)
	if err != nil {
		t.Fatalf("failed to parse config: %v", err)
	}
	// Set max connections to 1 to force reuse of the same connection slot
	config.MaxConns = 1

	pool, err := pgxpool.NewWithConfig(context.Background(), config)
	if err != nil {
		t.Fatalf("failed to create pool: %v", err)
	}
	defer pool.Close()

	connector := NewConnector(pool)
	db := sql.OpenDB(connector)
	defer db.Close()

	// 1. Begin a transaction with a cancellable context
	ctx, cancel := context.WithCancel(context.Background())
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		t.Fatalf("failed to begin transaction: %v", err)
	}

	// 2. Perform a write operation or simple query
	_, err = tx.ExecContext(ctx, "SELECT 1")
	if err != nil {
		t.Fatalf("failed to execute query in transaction: %v", err)
	}

	// 3. Cancel the context immediately before commit
	cancel()

	// 4. Call Commit() and assert that it returns a context cancellation error
	err = tx.Commit()
	if err == nil {
		t.Fatal("expected error on cancelled commit, got nil")
	}
	if !errors.Is(err, context.Canceled) && !errors.Is(err, sql.ErrTxDone) {
		t.Logf("commit returned error: %v", err)
	}

	// 5. Verify pool health by running a query on a new connection.
	// Since MaxConns = 1, this will reuse the connection slot.
	// If the previous connection was not properly discarded, this query will fail
	// or return transaction state errors.
	ctx2, cancel2 := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel2()

	var val int
	err = db.QueryRowContext(ctx2, "SELECT 1").Scan(&val)
	if err != nil {
		t.Fatalf("failed to execute query on reused connection: %v", err)
	}
	if val != 1 {
		t.Fatalf("expected 1, got %d", val)
	}
}