//go:build test // +build test package testdb import ( "context" "fmt" "log" "os" "path/filepath" "strings" "testing" "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" ) // Querier is a minimal interface satisfied by *db.PoolProxy, *pgxpool.Pool, // and pgx.Tx. Defined locally to avoid an import cycle (db → testdb → db). type Querier interface { Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) } const ( defaultTestDSN = "postgres://myuser:mypassword@localhost:5432/crussell_test?sslmode=disable" adminDatabaseDSN = "postgres://myuser:mypassword@localhost:5432/mydb?sslmode=disable" ) // CreateTestDatabase creates a fresh isolated test database, runs the schema migration, // and returns a pool connected to it. If a database with the same name exists, it's // dropped first. This enables parallel test execution across packages since each // package gets its own database running migrations in parallel. func CreateTestDatabase(dbName string) *pgxpool.Pool { ctx := context.Background() adminPool, err := pgxpool.New(ctx, adminDatabaseDSN) if err != nil { log.Fatalf("testdb: failed to connect to admin database: %v", err) } defer adminPool.Close() safeName := pgx.Identifier{dbName}.Sanitize() // Drop existing database with retries (terminate connections first) for attempt := 0; attempt < 3; attempt++ { _, _ = adminPool.Exec(ctx, fmt.Sprintf(` SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '%s' AND pid != pg_backend_pid() `, dbName)) _, err = adminPool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", safeName)) if err == nil { break } log.Printf("testdb: retry %d dropping %s: %v", attempt+1, dbName, err) } if err != nil { log.Fatalf("testdb: failed to drop database %s after retries: %v", dbName, err) } _, err = adminPool.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s", safeName)) if err != nil { log.Fatalf("testdb: failed to create database %s: %v", dbName, err) } adminPool.Close() testDSN := fmt.Sprintf("postgres://myuser:mypassword@localhost:5432/%s?sslmode=disable", dbName) poolCfg, err := pgxpool.ParseConfig(testDSN) if err != nil { log.Fatalf("testdb: failed to parse config: %v", err) } poolCfg.MaxConns = 16 poolCfg.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol pool, err := pgxpool.NewWithConfig(ctx, poolCfg) if err != nil { log.Fatalf("testdb: failed to connect to %s: %v", dbName, err) } if err := pool.Ping(ctx); err != nil { log.Fatalf("testdb: failed to ping %s: %v", dbName, err) } if err := migratePool(pool); err != nil { log.Fatalf("testdb: migration failed: %v", err) } return pool } // DestroyTestDatabase closes the pool (if non-nil) and drops the test database. // Retries aggressively with exponential backoff, pg_terminate_backend between // attempts, and a final DROP DATABASE WITH (FORCE) on PostgreSQL 13+. func DestroyTestDatabase(pool *pgxpool.Pool, dbName string) { if dbName == "" { return } if pool != nil { pool.Close() } ctx := context.Background() adminPool, err := pgxpool.New(ctx, adminDatabaseDSN) if err != nil { log.Printf("testdb: warning: failed to connect to drop %s: %v", dbName, err) return } defer adminPool.Close() safeName := pgx.Identifier{dbName}.Sanitize() // Aggressive retry with pg_terminate_backend between attempts and final FORCE option. for attempt := 0; attempt < 10; attempt++ { // Kill all connections to the target database. _, _ = adminPool.Exec(ctx, fmt.Sprintf(` SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '%s' AND pid != pg_backend_pid() `, dbName)) // Try DROP DATABASE WITH (FORCE) first (PG13+) — this kills remaining connections atomically. _, err = adminPool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", safeName)) if err == nil { return } // Fallback to plain DROP (may work if pg_terminate_backend was sufficient). _, err = adminPool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", safeName)) if err == nil { return } // Exponential backoff: 100ms, 200ms, 400ms, 800ms, 1.6s, 3.2s, 6.4s, 12.8s, 25.6s, 51.2s wait := time.Duration(100*(1< 80 { firstLine = firstLine[:80] + "..." } _, err := conn.Exec(ctx, stmt) if err != nil { return fmt.Errorf("failed to execute migration statement [%d]: %s\nError: %w", i+1, firstLine, err) } } return nil } func Migrate(t *testing.T, pool *pgxpool.Pool) { t.Helper() if err := migratePool(pool); err != nil { t.Fatal(err.Error()) } } func splitSQLStatements(sql string) []string { var statements []string var currentStatement strings.Builder inDollarBlock := false lines := strings.Split(sql, "\n") for _, line := range lines { trimmed := strings.TrimSpace(line) if trimmed == "" { continue } if strings.HasPrefix(trimmed, "--") { continue } currentStatement.WriteString(line) currentStatement.WriteString("\n") // Count occurrences of $$ in the line (handles single-line $$ blocks too) count := strings.Count(line, "$$") if count%2 == 1 { inDollarBlock = !inDollarBlock } if !inDollarBlock && strings.HasSuffix(trimmed, ";") { statements = append(statements, currentStatement.String()) currentStatement.Reset() } } if currentStatement.Len() > 0 { statements = append(statements, currentStatement.String()) } return statements } func Tx(t *testing.T, pool *pgxpool.Pool) pgx.Tx { t.Helper() ctx := context.Background() tx, err := pool.Begin(ctx) if err != nil { t.Fatalf("Failed to begin transaction: %v", err) } return tx } func TxWithRollback(t *testing.T, pool *pgxpool.Pool) (pgx.Tx, func()) { tx := Tx(t, pool) return tx, func() { tx.Rollback(context.Background()) } } func FindInitScript() (string, error) { cwd, err := os.Getwd() if err != nil { cwd = "" } paths := []string{ "../../../init-scripts/init-script.sql", "../../init-scripts/init-script.sql", "../init-scripts/init-script.sql", "init-scripts/init-script.sql", } if cwd != "" { paths = append(paths, filepath.Join(cwd, "..", "..", "init-scripts", "init-script.sql")) } for _, p := range paths { if _, err := os.Stat(p); err == nil { return p, nil } } return "", fmt.Errorf("could not find init-script.sql") }