Behavision/server/internal/migrate/migrate.go

// Package migrate applies the schema and records what it applied.
//
// Until this existed the migrations were run by hand - `psql < 001.sql` for
// each file, in order, by whoever remembered - and nothing anywhere recorded
// which had run. Three consequences, all of which had already happened:
//
//   - Re-running them against an existing database fails on the first CREATE
//     TABLE, so the setup script only worked once.
//   - Adding migration 008 to a release gave the operator no way to know
//     whether a given estate had it. The failure of a missed migration is not
//     a startup error; it is a query referencing a column that is not there,
//     surfacing on whichever endpoint touches it first.
//   - A half-applied migration - the file interrupted midway - left a schema
//     nothing could describe.
//
// So: every file runs inside ONE transaction together with the row that
// records it. Either both happen or neither does, and a database can always
// say which version it is on.
package migrate

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io/fs"
	"sort"
	"strconv"
	"strings"

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

// lockID namespaces the advisory lock this package takes. Two servers starting
// at the same moment - which is the normal shape of a rolling restart - must
// not both decide migration 008 is pending and both run it.
const lockID int64 = 7623094512340001

// File is one migration on disk.
type File struct {
	Version  int
	Name     string
	SQL      string
	Checksum string
}

// Record is one migration as the database remembers it.
type Record struct {
	Version   int
	Name      string
	Checksum  string
	Baselined bool
}

// Load reads and orders the migrations, rejecting anything it cannot place.
func Load(src fs.FS) ([]File, error) {
	entries, err := fs.Glob(src, "*.sql")
	if err != nil {
		return nil, err
	}
	seen := map[int]string{}
	out := make([]File, 0, len(entries))
	for _, name := range entries {
		version, err := versionOf(name)
		if err != nil {
			return nil, err
		}
		if other, dup := seen[version]; dup {
			// Two files claiming one version means the order they run in is
			// whatever the filesystem felt like, which is not an order.
			return nil, fmt.Errorf("migrations %s and %s share version %d",
				other, name, version)
		}
		seen[version] = name
		body, err := fs.ReadFile(src, name)
		if err != nil {
			return nil, err
		}
		sum := sha256.Sum256(body)
		out = append(out, File{
			Version:  version,
			Name:     name,
			SQL:      string(body),
			Checksum: hex.EncodeToString(sum[:]),
		})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Version < out[j].Version })
	return out, nil
}

func versionOf(name string) (int, error) {
	prefix, _, ok := strings.Cut(name, "_")
	if !ok {
		return 0, fmt.Errorf("migration %q is not named NNN_description.sql", name)
	}
	n, err := strconv.Atoi(prefix)
	if err != nil || n <= 0 {
		return 0, fmt.Errorf("migration %q does not start with a version number", name)
	}
	return n, nil
}

// Apply runs every migration the database has not recorded, in order.
//
// Returns the names it applied, so a caller can log "nothing to do" rather
// than staying silent about work it did not need to do.
func Apply(ctx context.Context, pool *pgxpool.Pool, src fs.FS) ([]string, error) {
	files, err := Load(src)
	if err != nil {
		return nil, err
	}

	conn, err := pool.Acquire(ctx)
	if err != nil {
		return nil, err
	}
	defer conn.Release()

	// Taken on this one connection and held for the whole run: a second server
	// booting waits here rather than racing us, and the lock is released with
	// the connection even if this process is killed.
	if _, err := conn.Exec(ctx, `SELECT pg_advisory_lock($1)`, lockID); err != nil {
		return nil, fmt.Errorf("migration lock: %w", err)
	}
	defer func() {
		_, _ = conn.Exec(context.WithoutCancel(ctx),
			`SELECT pg_advisory_unlock($1)`, lockID)
	}()

	if err := ensureTable(ctx, conn.Conn()); err != nil {
		return nil, err
	}
	done, err := recorded(ctx, conn.Conn())
	if err != nil {
		return nil, err
	}

	var applied []string
	for _, f := range files {
		prev, ok := done[f.Version]
		if ok {
			// An already-applied file that has since been edited. The database
			// does NOT contain what the repository says it does, and running
			// the new text now would apply half of it twice. Refuse and say so:
			// the fix is a new migration, never an edited one.
			if prev.Checksum != f.Checksum && !prev.Baselined {
				return applied, fmt.Errorf(
					"migration %s has changed since it was applied "+
						"(recorded %s, file %s) - add a new migration instead of editing one",
					f.Name, short(prev.Checksum), short(f.Checksum))
			}
			continue
		}
		if err := applyOne(ctx, conn.Conn(), f); err != nil {
			return applied, err
		}
		applied = append(applied, f.Name)
	}
	return applied, nil
}

// applyOne runs one file and records it in the SAME transaction. A migration
// that ran but was not recorded would run again on the next start; one that was
// recorded but did not run leaves a schema missing a column nothing will ever
// add.
func applyOne(ctx context.Context, conn *pgx.Conn, f File) error {
	tx, err := conn.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(context.WithoutCancel(ctx)) }()

	if _, err := tx.Exec(ctx, f.SQL); err != nil {
		return fmt.Errorf("migration %s: %w", f.Name, err)
	}
	if _, err := tx.Exec(ctx,
		`INSERT INTO schema_migrations (version, name, checksum) VALUES ($1, $2, $3)`,
		f.Version, f.Name, f.Checksum); err != nil {
		return fmt.Errorf("recording migration %s: %w", f.Name, err)
	}
	return tx.Commit(ctx)
}

// Baseline records migrations up to and including `through` as applied WITHOUT
// running them.
//
// For the one case that cannot be handled automatically: a database built
// before this package existed, whose schema is already there and whose history
// is not. Guessing is not an option - "the clients table exists" does not say
// whether migration 007's index does - so an operator states it, once, and it
// is recorded as a baseline rather than as a real application.
func Baseline(ctx context.Context, pool *pgxpool.Pool, src fs.FS, through int) ([]string, error) {
	files, err := Load(src)
	if err != nil {
		return nil, err
	}
	conn, err := pool.Acquire(ctx)
	if err != nil {
		return nil, err
	}
	defer conn.Release()
	if err := ensureTable(ctx, conn.Conn()); err != nil {
		return nil, err
	}
	var marked []string
	for _, f := range files {
		if f.Version > through {
			continue
		}
		tag, err := conn.Exec(ctx,
			`INSERT INTO schema_migrations (version, name, checksum, baselined)
			 VALUES ($1, $2, $3, true) ON CONFLICT (version) DO NOTHING`,
			f.Version, f.Name, f.Checksum)
		if err != nil {
			return marked, err
		}
		if tag.RowsAffected() == 1 {
			marked = append(marked, f.Name)
		}
	}
	return marked, nil
}

// Status reports what the database has recorded, oldest first.
func Status(ctx context.Context, pool *pgxpool.Pool) ([]Record, error) {
	conn, err := pool.Acquire(ctx)
	if err != nil {
		return nil, err
	}
	defer conn.Release()
	if err := ensureTable(ctx, conn.Conn()); err != nil {
		return nil, err
	}
	done, err := recorded(ctx, conn.Conn())
	if err != nil {
		return nil, err
	}
	out := make([]Record, 0, len(done))
	for _, r := range done {
		out = append(out, r)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Version < out[j].Version })
	return out, nil
}

func ensureTable(ctx context.Context, conn *pgx.Conn) error {
	_, err := conn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS schema_migrations (
			version    integer PRIMARY KEY,
			name       text NOT NULL,
			checksum   text NOT NULL,
			applied_at timestamptz NOT NULL DEFAULT now(),
			-- true when the row records a migration that was NOT run here,
			-- because the schema predates this table. Kept distinct so an
			-- adopted database never looks like one this code built.
			baselined  boolean NOT NULL DEFAULT false
		)`)
	return err
}

func recorded(ctx context.Context, conn *pgx.Conn) (map[int]Record, error) {
	rows, err := conn.Query(ctx,
		`SELECT version, name, checksum, baselined FROM schema_migrations`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[int]Record{}
	for rows.Next() {
		var r Record
		if err := rows.Scan(&r.Version, &r.Name, &r.Checksum, &r.Baselined); err != nil {
			return nil, err
		}
		out[r.Version] = r
	}
	return out, rows.Err()
}

func short(sum string) string {
	if len(sum) > 12 {
		return sum[:12]
	}
	return sum
}