Ensure transactions are closed in pqtest

Just calling db.Close() isn't enough, as it waits for queries to finish,
and it won't kill the BEGIN; etc.

For regular PostgreSQL tests this isn't a huge problem, most of the
time, but it does show up with something like "go test -count=50"
because eventually it will run out of connections.

e.g. pgbouncer and Supavisor tend to run out of connections even sooner,
with pgbouncer just waiting indefinitely until a connection gets closed,
causing the pgbouncer tests to be flaky on the CI.

Also wrap everything in TestMain() that tests if there are no
connections at the end of the test run.

Author: arp242

conn_test.go

@@ -1073,6 +1073,7 @@ func TestStmtQueryContext(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
+			defer stmt.Close()
 			_, err = stmt.QueryContext(ctx)
 			if !pqtest.ErrorContains(err, tt.wantErr) {
 				t.Errorf("wrong error:\nhave: %s\nwant: %s", err, tt.wantErr)
@@ -1116,6 +1117,7 @@ func TestStmtExecContext(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
+			defer stmt.Close()
 			_, err = stmt.ExecContext(ctx)
 			if !pqtest.ErrorContains(err, tt.wantErr) {
 				t.Errorf("wrong error:\nhave: %s\nwant: %s", err, tt.wantErr)
@@ -1745,7 +1747,6 @@ func BenchmarkSelect(b *testing.B) {
 func BenchmarkPreparedSelect(b *testing.B) {
 	run := func(b *testing.B, result any, query string) {
 		stmt := pqtest.Prepare(b, pqtest.MustDB(b), query).Stmt
-		defer stmt.Close()
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
@@ -1808,6 +1809,7 @@ func BenchmarkPreparedSelect(b *testing.B) {
 			if err != nil {
 				b.Fatal(err)
 			}
+			defer stmt.Close()
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
@@ -1828,6 +1830,7 @@ func BenchmarkPreparedSelect(b *testing.B) {
 			if err != nil {
 				b.Fatal(err)
 			}
+			defer stmt.Close()
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {

copy_test.go

@@ -45,10 +45,11 @@ func TestCopyInError(t *testing.T) {
 
 func TestCopyInErrorWrongType(t *testing.T) {
 	t.Parallel()
-	tx := pqtest.Begin(t, pqtest.MustDB(t))
+	db := pqtest.MustDB(t)
+	tx := pqtest.Begin(t, db)
 	pqtest.Exec(t, tx, `create temp table tbl (num integer)`)
 
-	stmt := pqtest.Prepare(t, tx, `copy tbl (num) from stdin`)
+	stmt := pqtest.Prepare(t, tx, `copy tbl (num) from stdin`, db)
 	stmt.MustExec(t, "Héllö\n ☃!\r\t\\")
 	_, err := stmt.Exec()
 	mustAs(t, err, pqerror.InvalidTextRepresentation)
@@ -66,10 +67,11 @@ func TestCopyInErrorOutsideTransaction(t *testing.T) {
 
 func TestCopyInQueryWhileCopy(t *testing.T) {
 	t.Parallel()
-	tx := pqtest.Begin(t, pqtest.MustDB(t))
+	db := pqtest.MustDB(t)
+	tx := pqtest.Begin(t, db)
 	pqtest.Exec(t, tx, `create temp table tbl (i int primary key)`)
 
-	pqtest.Prepare(t, tx, "copy tbl (i) from stdin")
+	pqtest.Prepare(t, tx, "copy tbl (i) from stdin", db)
 	_, err := tx.Query(`select 1`)
 	if !errors.Is(err, errQueryInProgress) {
 		t.Errorf("wrong error:\nhave: %s\nwant: %s", err, errQueryInProgress)
@@ -97,10 +99,11 @@ func TestCopyInNull(t *testing.T) {
 		tt := tt
 		t.Run("", func(t *testing.T) {
 			t.Parallel()
-			tx := pqtest.Begin(t, pqtest.MustDB(t))
+			db := pqtest.MustDB(t)
+			tx := pqtest.Begin(t, db)
 
 			pqtest.Exec(t, tx, `create temp table tbl (i int, t text)`)
-			stmt := pqtest.Prepare(t, tx, tt.copy)
+			stmt := pqtest.Prepare(t, tx, tt.copy, db)
 			stmt.MustExec(t, 42, "forty-two")
 			stmt.MustExec(t, tt.null, tt.null)
 			stmt.MustExec(t)
@@ -130,10 +133,11 @@ func TestCopyInMultipleValues(t *testing.T) {
 		tt := tt
 		t.Run("", func(t *testing.T) {
 			t.Parallel()
-			tx := pqtest.Begin(t, pqtest.MustDB(t))
+			db := pqtest.MustDB(t)
+			tx := pqtest.Begin(t, db)
 			pqtest.Exec(t, tx, `create temp table tbl (a int, b varchar)`)
 
-			stmt := pqtest.Prepare(t, tx, tt.query)
+			stmt := pqtest.Prepare(t, tx, tt.query, db)
 			for i := 0; i < 500; i++ {
 				stmt.MustExec(t, int64(i), strings.Repeat("#", 500))
 			}
@@ -161,7 +165,8 @@ func TestCopyInMultipleValues(t *testing.T) {
 
 func TestCopyInRaiseStmtTrigger(t *testing.T) {
 	t.Parallel()
-	tx := pqtest.Begin(t, pqtest.MustDB(t))
+	db := pqtest.MustDB(t)
+	tx := pqtest.Begin(t, db)
 	pqtest.Exec(t, tx, `create temp table tbl (a int, b varchar)`)
 	pqtest.Exec(t, tx, `
 		create or replace function pg_temp.temptest()
@@ -178,7 +183,7 @@ func TestCopyInRaiseStmtTrigger(t *testing.T) {
 		for each row execute procedure pg_temp.temptest()
 	`)
 
-	stmt := pqtest.Prepare(t, tx, `copy tbl (a, b) from stdin`)
+	stmt := pqtest.Prepare(t, tx, `copy tbl (a, b) from stdin`, db)
 	stmt.MustExec(t, int64(1), strings.Repeat("#", 500))
 	stmt.MustExec(t)
 	stmt.MustClose(t)
@@ -195,10 +200,11 @@ func TestCopyInRaiseStmtTrigger(t *testing.T) {
 
 func TestCopyInTypes(t *testing.T) {
 	t.Parallel()
-	tx := pqtest.Begin(t, pqtest.MustDB(t))
+	db := pqtest.MustDB(t)
+	tx := pqtest.Begin(t, db)
 	pqtest.Exec(t, tx, `create temp table tbl (num integer, text varchar, blob bytea, nothing varchar)`)
 
-	stmt := pqtest.Prepare(t, tx, `copy tbl (num, text, blob, nothing) from stdin`)
+	stmt := pqtest.Prepare(t, tx, `copy tbl (num, text, blob, nothing) from stdin`, db)
 	stmt.MustExec(t, int64(1234567890), "Héllö\n ☃!\r\t\\", []byte{0, 255, 9, 10, 13}, nil)
 	stmt.MustExec(t)
 	stmt.MustClose(t)
@@ -241,7 +247,7 @@ func TestCopyInRespLoopConnectionError(t *testing.T) {
 
 	pid := pqtest.Query[int64](t, tx, `select pg_backend_pid() as pid`)
 	pqtest.Exec(t, tx, "create temp table tbl (a int)")
-	stmt := pqtest.Prepare(t, tx, `copy tbl (a) from stdin`)
+	stmt := pqtest.Prepare(t, tx, `copy tbl (a) from stdin`, db)
 	pqtest.Exec(t, db, `select pg_terminate_backend($1)`, pid[0]["pid"])
 
 	var err error
@@ -271,10 +277,11 @@ func TestCopyInRespLoopConnectionError(t *testing.T) {
 }
 
 func BenchmarkCopyIn(b *testing.B) {
-	tx := pqtest.Begin(b, pqtest.MustDB(b))
+	db := pqtest.MustDB(b)
+	tx := pqtest.Begin(b, db)
 
 	pqtest.Exec(b, tx, `create temp table tbl (a int, b varchar)`)
-	stmt := pqtest.Prepare(b, tx, `copy tbl (a, b) from stdin`)
+	stmt := pqtest.Prepare(b, tx, `copy tbl (a, b) from stdin`, db)
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {

encode_test.go

@@ -334,7 +334,7 @@ func TestByteaOutputFormats(t *testing.T) {
 			}
 
 			{ // Same but with Prepare
-				stmt := pqtest.Prepare(t, tx, `select decode('5c7800ff6162630108', 'hex')`)
+				stmt := pqtest.Prepare(t, tx, `select decode('5c7800ff6162630108', 'hex')`, db)
 				rows, err := stmt.Query()
 				if err != nil {
 					t.Fatal(err)

error_test.go

@@ -255,6 +255,7 @@ func TestNetworkError(t *testing.T) {
 	}
 	c.Dialer(failDialer{})
 	db := sql.OpenDB(c)
+	defer db.Close()
 	db.SetMaxIdleConns(1)
 	db.SetMaxOpenConns(1)
 	if err := db.Ping(); err != nil {

example_test.go

@@ -142,6 +142,7 @@ func Example_copyFromStdin() {
 	if err != nil {
 		log.Fatal(err)
 	}
+	defer stmt.Close()
 
 	// Insert rows.
 	users := []struct {

helper_test.go

@@ -102,6 +102,11 @@ func Home(t *testing.T) string {
 	return pqutil.Home(true)
 }
 
+var (
+	mu       sync.Mutex
+	cleanups = make(map[*sql.DB][]func())
+)
+
 // DB connects to the test database and returns the Ping error. The connection
 // is closed in t.Cleanup().
 func DB(t testing.TB, conninfo ...string) (*sql.DB, error) {
@@ -110,7 +115,15 @@ func DB(t testing.TB, conninfo ...string) (*sql.DB, error) {
 	if err != nil {
 		t.Fatalf("pqtest.DB: %s", err)
 	}
-	t.Cleanup(func() { db.Close() })
+	t.Cleanup(func() {
+		mu.Lock()
+		defer mu.Unlock()
+		for i := len(cleanups[db]) - 1; i >= 0; i-- {
+			cleanups[db][i]()
+		}
+		delete(cleanups, db)
+		db.Close()
+	})
 	return db, db.Ping()
 }
 
@@ -132,10 +145,9 @@ func Begin(t testing.TB, db *sql.DB) *sql.Tx {
 	if err != nil {
 		t.Fatalf("pqtest.Begin: %s", err)
 	}
-	// We can't call t.Cleanup here as that will race with the t.Cleanup from
-	// MustDB (it's called in "last added, first called", so the tx.Rollback
-	// gets called after db.Close)
-	// t.Cleanup(func() { tx.Rollback() })
+	mu.Lock()
+	defer mu.Unlock()
+	cleanups[db] = append(cleanups[db], func() { tx.Rollback() })
 	return tx
 }
 
@@ -163,12 +175,25 @@ func (s *Stmt) MustClose(t testing.TB) {
 // Prepare a new statement, calling t.Fatal() if this fails.
 func Prepare(t testing.TB, db interface {
 	Prepare(string) (*sql.Stmt, error)
-}, q string) *Stmt {
+}, q string, sqldb ...*sql.DB) *Stmt {
 	t.Helper()
 	stmt, err := db.Prepare(q)
 	if err != nil {
 		t.Fatalf("pqtest.Prepare: %s", err)
 	}
+
+	if len(sqldb) == 0 {
+		sqldb = make([]*sql.DB, 1)
+		var ok bool
+		sqldb[0], ok = db.(*sql.DB)
+		if !ok {
+			t.Fatalf("pqtest.Prepare: must pass sql.DB when using transaction")
+		}
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	cleanups[sqldb[0]] = append(cleanups[sqldb[0]], func() { stmt.Close() })
+
 	return &Stmt{stmt}
 }
 

internal/pqtest/pqtest.go

@@ -0,0 +1,33 @@
+package pqtest_test
+
+import (
+	"testing"
+
+	_ "github.com/lib/pq"
+	"github.com/lib/pq/internal/pqtest"
+)
+
+// Just calling db.Close() isn't enough, as it waits for queries to finish, and
+// it won't kill the BEGIN; etc.
+func TestCleanup(t *testing.T) {
+	t.Run("", func(t *testing.T) {
+		t.Setenv("PGAPPNAME", "pqgo-cleanup")
+		db := pqtest.MustDB(t)
+		pqtest.Begin(t, db)
+		pqtest.Query[int](t, db, `select 1`)
+		stmt := pqtest.Prepare(t, db, `select 1`)
+
+		// No helper function for these as they're not used that frequently, and
+		// for stmt.Query() also difficult to do right.
+		rows1, _ := db.Query(`select 1`)
+		defer rows1.Close()
+		rows2, _ := stmt.Query()
+		defer rows2.Close()
+	})
+
+	rows := pqtest.Query[any](t, pqtest.MustDB(t),
+		`select pid, query from pg_stat_activity where application_name = 'pqgo-cleanup' and pid != pg_backend_pid()`)
+	for _, r := range rows {
+		t.Errorf("connection still active: pid=%d; query=%q\n", r["pid"], r["query"])
+	}
+}

internal/pqtest/pqtest_test.go

@@ -9,7 +9,7 @@ import (
 	"math/big"
 	"net"
 	"runtime"
-	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -181,24 +181,23 @@ func TestListenerConnExecDeadlock(t *testing.T) {
 	l, _ := newTestListenerConn(t)
 	defer l.Close()
 
-	var wg sync.WaitGroup
-	wg.Add(2)
+	var done atomic.Int32
 	go func() {
 		l.ExecSimpleQuery("select pg_sleep(0.2)")
-		wg.Done()
+		done.Add(1)
 	}()
 	runtime.Gosched()
 	go func() {
 		l.ExecSimpleQuery("select 1")
-		wg.Done()
+		done.Add(1)
 	}()
-	// give the two goroutines some time to get into position
-	runtime.Gosched()
-	// calls Close on the net.Conn; equivalent to a network failure
-	l.Close()
+	runtime.Gosched() // Give the above goroutine some time to get into position.
+	l.Close()         // Calls Close on the net.Conn; equivalent to a network failure.
 
-	defer time.AfterFunc(200*time.Millisecond, func() { panic("timed out") }).Stop()
-	wg.Wait()
+	time.Sleep(200 * time.Millisecond)
+	if done.Load() != 2 {
+		t.Fatal("timed out")
+	}
 }
 
 // Test for ListenerConn being closed while a slow query is executing
@@ -207,29 +206,26 @@ func TestListenerConnCloseWhileQueryIsExecuting(t *testing.T) {
 	l, _ := newTestListenerConn(t)
 	defer l.Close()
 
-	var wg sync.WaitGroup
-	wg.Add(1)
-
+	var done atomic.Int32
 	go func() {
 		sent, err := l.ExecSimpleQuery("select pg_sleep(0.2)")
 		if sent {
 			panic("expected sent=false")
 		}
-		// could be any of a number of errors
-		if err == nil {
+		if err == nil { // Could be any of a number of errors.
 			panic("expected error")
 		}
-		wg.Done()
+		done.Add(1)
 	}()
-	// give the above goroutine some time to get into position
-	runtime.Gosched()
-	err := l.Close()
-	if err != nil {
+	runtime.Gosched() // Give the above goroutine some time to get into position.
+	if err := l.Close(); err != nil {
 		t.Fatal(err)
 	}
 
-	defer time.AfterFunc(200*time.Millisecond, func() { panic("timed out") }).Stop()
-	wg.Wait()
+	time.Sleep(200 * time.Millisecond)
+	if done.Load() != 1 {
+		t.Fatal("timed out")
+	}
 }
 
 func TestListenerNotifyExtra(t *testing.T) {