EvoBGP/internal/db/postgres_integration_test.go

package db

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"io/fs"
	"os"
	"strings"
	"sync"
	"testing"
	"testing/fstest"
	"time"

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

func TestPostgresMigrationsIntegration(t *testing.T) {
	pool := isolatedMigrationPool(t)
	ctx := context.Background()
	if err := ApplyPostgresMigrations(ctx, pool, migrations.Postgres, "postgres"); err != nil {
		t.Fatalf("first migration pass: %v", err)
	}
	if err := ApplyPostgresMigrations(ctx, pool, migrations.Postgres, "postgres"); err != nil {
		t.Fatalf("second migration pass: %v", err)
	}
	files, err := fs.ReadDir(migrations.Postgres, "postgres")
	if err != nil {
		t.Fatal(err)
	}
	want := 0
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".up.sql") {
			want++
		}
	}
	var got int
	if err := pool.QueryRow(ctx, `SELECT count(*)::int FROM schema_migrations`).Scan(&got); err != nil {
		t.Fatal(err)
	}
	if got != want {
		t.Fatalf("applied migrations = %d, want %d", got, want)
	}
	var extension string
	if err := pool.QueryRow(ctx, `SELECT extname FROM pg_extension WHERE extname = 'pg_stat_statements'`).Scan(&extension); err != nil {
		t.Fatalf("pg_stat_statements extension missing: %v", err)
	}
}

func TestPostgresMigrationConcurrentApplyIntegration(t *testing.T) {
	pool := isolatedMigrationPool(t)
	files := fstest.MapFS{"sql/000001_probe.up.sql": {Data: []byte(`CREATE TABLE migration_probe (id integer PRIMARY KEY); SELECT pg_sleep(0.1);`)}}
	var wg sync.WaitGroup
	errors := make(chan error, 2)
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			errors <- ApplyPostgresMigrations(context.Background(), pool, files, "sql")
		}()
	}
	wg.Wait()
	close(errors)
	for err := range errors {
		if err != nil {
			t.Fatal(err)
		}
	}
	assertMigrationVersionCount(t, pool, 1)
}

func TestPostgresMigrationRollbackIntegration(t *testing.T) {
	pool := isolatedMigrationPool(t)
	files := fstest.MapFS{"sql/000001_broken.up.sql": {Data: []byte(`CREATE TABLE migration_partial (id integer); SELECT 1 / 0;`)}}
	if err := ApplyPostgresMigrations(context.Background(), pool, files, "sql"); err == nil {
		t.Fatal("expected migration error")
	}
	assertMigrationVersionCount(t, pool, 0)
	var exists bool
	if err := pool.QueryRow(context.Background(), `SELECT to_regclass('migration_partial') IS NOT NULL`).Scan(&exists); err != nil {
		t.Fatal(err)
	}
	if exists {
		t.Fatal("failed migration left a table behind")
	}
	// A transaction terminator hidden behind a comment must be rejected before
	// the first DDL statement can run.
	files = fstest.MapFS{"sql/000001_broken.up.sql": {Data: []byte(`CREATE TABLE migration_partial (id integer); /* boundary */ COMMIT; SELECT 1 / 0;`)}}
	if err := ApplyPostgresMigrations(context.Background(), pool, files, "sql"); err == nil || !strings.Contains(err.Error(), "unsupported non-transactional SQL") {
		t.Fatalf("comment-hidden COMMIT error = %v", err)
	}
	assertMigrationVersionCount(t, pool, 0)
	if err := pool.QueryRow(context.Background(), `SELECT to_regclass('migration_partial') IS NOT NULL`).Scan(&exists); err != nil {
		t.Fatal(err)
	}
	if exists {
		t.Fatal("rejected migration left a table behind")
	}
	files = fstest.MapFS{"sql/000001_broken.up.sql": {Data: []byte(`CREATE TABLE foo$tag$ (id int); COMMIT; -- $tag$`)}}
	if err := ApplyPostgresMigrations(context.Background(), pool, files, "sql"); err == nil || !strings.Contains(err.Error(), "unsupported non-transactional SQL") {
		t.Fatalf("identifier-embedded dollar tag error = %v", err)
	}
	assertMigrationVersionCount(t, pool, 0)
	if err := pool.QueryRow(context.Background(), `SELECT to_regclass('foo$tag$') IS NOT NULL`).Scan(&exists); err != nil {
		t.Fatal(err)
	}
	if exists {
		t.Fatal("identifier-embedded dollar tag left a table behind")
	}
}

func TestPostgresMigrationIdempotentIntegration(t *testing.T) {
	pool := isolatedMigrationPool(t)
	files := fstest.MapFS{"sql/000001_once.up.sql": {Data: []byte(`CREATE TABLE migration_once (id integer); INSERT INTO migration_once VALUES (1);`)}}
	for i := 0; i < 2; i++ {
		if err := ApplyPostgresMigrations(context.Background(), pool, files, "sql"); err != nil {
			t.Fatal(err)
		}
	}
	assertMigrationVersionCount(t, pool, 1)
	var rows int
	if err := pool.QueryRow(context.Background(), `SELECT count(*) FROM migration_once`).Scan(&rows); err != nil {
		t.Fatal(err)
	}
	if rows != 1 {
		t.Fatalf("migration rows = %d, want 1", rows)
	}
}

func TestPostgresMigrationLockCancellationIntegration(t *testing.T) {
	pool := isolatedMigrationPool(t)
	locker, err := pool.Acquire(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer locker.Release()
	if _, err := locker.Exec(context.Background(), `SELECT pg_advisory_lock($1)`, postgresMigrationLockKey); err != nil {
		t.Fatal(err)
	}
	defer func() {
		_, _ = locker.Exec(context.Background(), `SELECT pg_advisory_unlock($1)`, postgresMigrationLockKey)
	}()
	files := fstest.MapFS{"sql/000001_probe.up.sql": {Data: []byte(`CREATE TABLE lock_probe (id integer);`)}}
	ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond)
	defer cancel()
	err = ApplyPostgresMigrations(ctx, pool, files, "sql")
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("lock wait error = %v, want context deadline exceeded", err)
	}
	var exists bool
	if err := pool.QueryRow(context.Background(), `SELECT to_regclass('schema_migrations') IS NOT NULL`).Scan(&exists); err != nil {
		t.Fatal(err)
	}
	if exists {
		t.Fatal("schema_migrations was created before migration lock acquisition")
	}
	var unlocked bool
	if err := locker.QueryRow(context.Background(), `SELECT pg_advisory_unlock($1)`, postgresMigrationLockKey).Scan(&unlocked); err != nil || !unlocked {
		t.Fatalf("release held lock: unlocked=%v err=%v", unlocked, err)
	}
	if err := ApplyPostgresMigrations(context.Background(), pool, files, "sql"); err != nil {
		t.Fatalf("apply after lock release: %v", err)
	}
}

func isolatedMigrationPool(t *testing.T) *pgxpool.Pool {
	t.Helper()
	dsn := os.Getenv("EVOBGP_TEST_DATABASE_URL")
	if dsn == "" {
		t.Skip("EVOBGP_TEST_DATABASE_URL not set")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	admin, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Fatal(err)
	}
	var nonce [8]byte
	if _, err := rand.Read(nonce[:]); err != nil {
		admin.Close()
		t.Fatal(err)
	}
	schema := "evobgp_mig_test_" + hex.EncodeToString(nonce[:])
	if _, err := admin.Exec(ctx, `CREATE SCHEMA `+schema); err != nil {
		admin.Close()
		t.Fatal(err)
	}
	var pool *pgxpool.Pool
	t.Cleanup(func() {
		if pool != nil {
			pool.Close()
		}
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cleanupCancel()
		if _, err := admin.Exec(cleanupCtx, `DROP SCHEMA `+schema+` CASCADE`); err != nil {
			t.Errorf("drop isolated schema %s: %v", schema, err)
		}
		admin.Close()
	})
	cfg, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.ConnConfig.RuntimeParams == nil {
		cfg.ConnConfig.RuntimeParams = make(map[string]string)
	}
	cfg.ConnConfig.RuntimeParams["search_path"] = schema
	cfg.MaxConns = 3
	pool, err = pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		t.Fatal(err)
	}
	return pool
}

func assertMigrationVersionCount(t *testing.T, pool *pgxpool.Pool, want int) {
	t.Helper()
	var got int
	if err := pool.QueryRow(context.Background(), `SELECT count(*) FROM schema_migrations`).Scan(&got); err != nil {
		t.Fatal(err)
	}
	if got != want {
		t.Fatalf("migration versions = %d, want %d", got, want)
	}
}