golib/pg/pgmigrator_test.go

package pg

import (
	"context"
	"errors"
	"reflect"
	"strings"
	"testing"
)

const testLockName = "golib-test-migrations"

func newTestMigrator(s *fakeSession) (*pgMigrator, *int) {
	discards := 0
	m := &pgMigrator{
		conn:    s,
		key:     LockKey(testLockName),
		discard: func(context.Context) { discards++ },
	}
	return m, &discards
}

// The lock and unlock must name the same key, and it must be the derived one:
// a mismatch here is a deadlock or a lock nobody else respects.
func TestPgMigrator_LockUnlockUseTheDerivedKey(t *testing.T) {
	s := &fakeSession{}
	m, _ := newTestMigrator(s)
	ctx := context.Background()

	if err := m.Lock(ctx); err != nil {
		t.Fatalf("Lock: %v", err)
	}
	if err := m.Unlock(ctx); err != nil {
		t.Fatalf("Unlock: %v", err)
	}

	want := []string{"SELECT pg_advisory_lock($1)", "SELECT pg_advisory_unlock($1)"}
	if got := s.execSQL(); !reflect.DeepEqual(got, want) {
		t.Fatalf("statements %v, want %v", got, want)
	}
	key := LockKey(testLockName)
	for _, c := range s.execs {
		if len(c.args) != 1 || c.args[0] != key {
			t.Fatalf("%q got args %v, want [%d]", c.sql, c.args, key)
		}
	}
}

func TestPgMigrator_LockAndUnlockPropagateErrors(t *testing.T) {
	lockErr := errors.New("lock failed")
	unlockErr := errors.New("unlock failed")
	s := &fakeSession{execErrs: map[string]error{
		"SELECT pg_advisory_lock($1)":   lockErr,
		"SELECT pg_advisory_unlock($1)": unlockErr,
	}}
	m, _ := newTestMigrator(s)
	if err := m.Lock(context.Background()); !errors.Is(err, lockErr) {
		t.Errorf("Lock error = %v, want %v", err, lockErr)
	}
	if err := m.Unlock(context.Background()); !errors.Is(err, unlockErr) {
		t.Errorf("Unlock error = %v, want %v", err, unlockErr)
	}
}

func TestPgMigrator_DiscardClosesTheSession(t *testing.T) {
	m, discards := newTestMigrator(&fakeSession{})
	m.Discard(context.Background())
	if *discards != 1 {
		t.Fatalf("discard called %d times, want 1", *discards)
	}
}

func TestPgMigrator_EnsureVersionTableRunsTheDDL(t *testing.T) {
	s := &fakeSession{}
	m, _ := newTestMigrator(s)
	if err := m.EnsureVersionTable(context.Background()); err != nil {
		t.Fatalf("EnsureVersionTable: %v", err)
	}
	if got := s.execSQL(); !reflect.DeepEqual(got, []string{schemaMigrationsDDL}) {
		t.Fatalf("statements %v, want the schema_migrations DDL", got)
	}
	// Re-running the DDL must be harmless, so it has to be IF NOT EXISTS.
	if !strings.Contains(schemaMigrationsDDL, "IF NOT EXISTS") {
		t.Error("schema_migrations DDL is not IF NOT EXISTS-guarded")
	}
}

func TestPgMigrator_EnsureVersionTablePropagatesErrors(t *testing.T) {
	want := errors.New("permission denied")
	s := &fakeSession{execErrs: map[string]error{schemaMigrationsDDL: want}}
	m, _ := newTestMigrator(s)
	if err := m.EnsureVersionTable(context.Background()); !errors.Is(err, want) {
		t.Fatalf("error = %v, want %v", err, want)
	}
}

func TestPgMigrator_AppliedVersionsReadsAndClosesRows(t *testing.T) {
	rows := &fakeRows{values: []string{"0001_first.sql", "0002_second.sql"}}
	s := &fakeSession{rows: rows}
	m, _ := newTestMigrator(s)

	got, err := m.AppliedVersions(context.Background())
	if err != nil {
		t.Fatalf("AppliedVersions: %v", err)
	}
	want := map[string]bool{"0001_first.sql": true, "0002_second.sql": true}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("versions %v, want %v", got, want)
	}
	if !rows.closed {
		t.Error("rows were not closed")
	}
	if wantQ := []string{"SELECT version FROM schema_migrations"}; !reflect.DeepEqual(s.queries, wantQ) {
		t.Fatalf("queries %v, want %v", s.queries, wantQ)
	}
}

func TestPgMigrator_AppliedVersionsEmptyTable(t *testing.T) {
	s := &fakeSession{rows: &fakeRows{}}
	m, _ := newTestMigrator(s)
	got, err := m.AppliedVersions(context.Background())
	if err != nil {
		t.Fatalf("AppliedVersions: %v", err)
	}
	if len(got) != 0 {
		t.Fatalf("versions %v, want an empty set", got)
	}
}

func TestPgMigrator_AppliedVersionsErrors(t *testing.T) {
	queryErr := errors.New("relation does not exist")
	scanErr := errors.New("bad column type")
	rowsErr := errors.New("connection reset mid-read")

	tests := []struct {
		name string
		sess *fakeSession
		want error
	}{
		{"query fails", &fakeSession{queryErr: queryErr}, queryErr},
		{"scan fails", &fakeSession{rows: &fakeRows{values: []string{"x"}, scanErr: scanErr}}, scanErr},
		// A read that dies part-way must not be reported as a complete set, or
		// already-applied migrations would be re-run.
		{"rows.Err after iteration", &fakeSession{rows: &fakeRows{values: []string{"x"}, err: rowsErr}}, rowsErr},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			m, _ := newTestMigrator(tc.sess)
			got, err := m.AppliedVersions(context.Background())
			if !errors.Is(err, tc.want) {
				t.Fatalf("error = %v, want %v", err, tc.want)
			}
			if got != nil {
				t.Fatalf("versions = %v, want nil on error", got)
			}
		})
	}
}

// One transaction carries both the migration body and its tracking row, so a
// crash can never leave the schema changed but unrecorded.
func TestPgMigrator_ApplyRunsBodyAndVersionRowInOneTx(t *testing.T) {
	const body = "CREATE TABLE IF NOT EXISTS a ();"
	tx := &fakeTx{}
	s := &fakeSession{tx: tx}
	m, _ := newTestMigrator(s)

	if err := m.Apply(context.Background(), "0001_first.sql", body); err != nil {
		t.Fatalf("Apply: %v", err)
	}
	want := []string{body, "INSERT INTO schema_migrations (version) VALUES ($1)"}
	if got := tx.execSQL(); !reflect.DeepEqual(got, want) {
		t.Fatalf("statements %v, want %v", got, want)
	}
	// The body must be sent with no arguments so pgx uses the simple protocol
	// and a multi-statement file runs.
	if len(tx.execs[0].args) != 0 {
		t.Errorf("migration body sent with args %v, want none", tx.execs[0].args)
	}
	if got := tx.execs[1].args; len(got) != 1 || got[0] != "0001_first.sql" {
		t.Errorf("version row args %v, want [0001_first.sql]", got)
	}
	if !tx.committed {
		t.Error("transaction was not committed")
	}
	if tx.rolled != 1 {
		t.Errorf("Rollback called %d times, want 1 (the no-op deferred rollback)", tx.rolled)
	}
}

func TestPgMigrator_ApplyRollsBackOnFailure(t *testing.T) {
	const body = "CREATE TABLE oops ();"
	bodyErr := errors.New("syntax error")
	insertErr := errors.New("duplicate key")
	const insert = "INSERT INTO schema_migrations (version) VALUES ($1)"

	tests := []struct {
		name      string
		execErrs  map[string]error
		commitErr error
		want      error
		wantExecs int
	}{
		{"body fails", map[string]error{body: bodyErr}, nil, bodyErr, 1},
		{"version row fails", map[string]error{insert: insertErr}, nil, insertErr, 2},
		{"commit fails", nil, errors.New("commit failed"), nil, 2},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			tx := &fakeTx{execErrs: tc.execErrs, commitErr: tc.commitErr}
			s := &fakeSession{tx: tx}
			m, _ := newTestMigrator(s)

			err := m.Apply(context.Background(), "0001_first.sql", body)
			if err == nil {
				t.Fatal("expected Apply to fail")
			}
			if tc.want != nil && !errors.Is(err, tc.want) {
				t.Fatalf("error = %v, want %v", err, tc.want)
			}
			if len(tx.execs) != tc.wantExecs {
				t.Errorf("%d statements ran, want %d", len(tx.execs), tc.wantExecs)
			}
			if tx.committed {
				t.Error("transaction was committed despite the failure")
			}
			if tx.rolled == 0 {
				t.Error("transaction was not rolled back")
			}
		})
	}
}

func TestPgMigrator_ApplyPropagatesBeginFailure(t *testing.T) {
	want := errors.New("cannot start transaction")
	m, _ := newTestMigrator(&fakeSession{beginErr: want})
	if err := m.Apply(context.Background(), "0001_first.sql", "SELECT 1;"); !errors.Is(err, want) {
		t.Fatalf("error = %v, want %v", err, want)
	}
}

// migrateSession is the whole runner over a real pgMigrator: lock, DDL, read,
// apply each file, unlock — in that order, on one session.
func TestMigrateSession_FullRunOnOneSession(t *testing.T) {
	tx := &fakeTx{}
	s := &fakeSession{rows: &fakeRows{}, tx: tx}
	discards := 0
	err := migrateSession(context.Background(), s,
		func(context.Context) { discards++ },
		testFS(),
		MigrateOptions{LockName: testLockName})
	if err != nil {
		t.Fatalf("migrateSession: %v", err)
	}

	want := []string{
		"SELECT pg_advisory_lock($1)",
		schemaMigrationsDDL,
		"SELECT pg_advisory_unlock($1)",
	}
	if got := s.execSQL(); !reflect.DeepEqual(got, want) {
		t.Fatalf("session statements %v, want %v", got, want)
	}
	// Every file's body plus its version row, all through the one fake tx.
	if got, wantN := len(tx.execs), 2*len(testNames); got != wantN {
		t.Fatalf("%d statements in transactions, want %d", got, wantN)
	}
	if discards != 0 {
		t.Errorf("connection discarded %d times after a clean run, want 0", discards)
	}
}

func TestMigrateSession_WrapsRunnerErrors(t *testing.T) {
	s := &fakeSession{execErrs: map[string]error{"SELECT pg_advisory_lock($1)": errors.New("no lock")}}
	err := migrateSession(context.Background(), s, func(context.Context) {}, testFS(),
		MigrateOptions{LockName: testLockName})
	if err == nil {
		t.Fatal("expected an error")
	}
	if !strings.Contains(err.Error(), "migrate schema") {
		t.Fatalf("error %q is not wrapped with the migration context", err)
	}
}

// An unlock that fails must reach the discard func through the real migrator,
// not just the fake one the runner tests use.
func TestMigrateSession_UnlockFailureDiscards(t *testing.T) {
	s := &fakeSession{
		rows:     &fakeRows{},
		tx:       &fakeTx{},
		execErrs: map[string]error{"SELECT pg_advisory_unlock($1)": errors.New("gone")},
	}
	discards := 0
	err := migrateSession(context.Background(), s, func(context.Context) { discards++ },
		testFS(), MigrateOptions{LockName: testLockName})
	if err != nil {
		t.Fatalf("migrateSession: %v", err)
	}
	if discards != 1 {
		t.Fatalf("connection discarded %d times, want 1", discards)
	}
}