Warp/server/internal/storage/migrations/migrations.go

// Package migrations runs embedded SQL migrations against any
// database/sql driver.
//
// Migrations live next to their driver as files named NNNN_description.sql
// and are applied in version order inside a transaction each, recorded in a
// schema_migrations table. A migration that has already run and whose file
// has since changed is an error, not a silent no-op.
package migrations

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"io/fs"
	"log/slog"
	"path"
	"slices"
	"strconv"
	"strings"
	"time"
)

// Migration is one SQL file.
type Migration struct {
	Version  int
	Name     string
	SQL      string
	Checksum string
}

// Applied is a row of schema_migrations.
type Applied struct {
	Version   int
	Name      string
	Checksum  string
	AppliedAt time.Time
}

// Dialect carries the handful of statements that differ per database. Each
// driver package supplies one; the runner itself stays dialect-agnostic.
type Dialect struct {
	// Name is used in log lines and errors, e.g. "sqlite".
	Name string

	// CreateTable creates schema_migrations if it does not exist, with
	// columns (version, name, checksum, applied_at).
	CreateTable string

	// SelectApplied reads every row, ordered by version.
	SelectApplied string

	// InsertApplied records one migration. Its placeholders bind
	// (version, name, checksum, applied_at) in that order.
	InsertApplied string
}

// Load reads every NNNN_*.sql file in dir of fsys.
func Load(fsys fs.FS, dir string) ([]Migration, error) {
	entries, err := fs.ReadDir(fsys, dir)
	if err != nil {
		return nil, fmt.Errorf("read migrations directory %s: %w", dir, err)
	}

	var out []Migration
	seen := map[int]string{}
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".sql") {
			continue
		}
		version, name, err := parseFilename(e.Name())
		if err != nil {
			return nil, err
		}
		if prev, dup := seen[version]; dup {
			return nil, fmt.Errorf("duplicate migration version %d: %s and %s", version, prev, e.Name())
		}
		seen[version] = e.Name()

		body, err := fs.ReadFile(fsys, path.Join(dir, e.Name()))
		if err != nil {
			return nil, fmt.Errorf("read migration %s: %w", e.Name(), err)
		}
		sum := sha256.Sum256(body)
		out = append(out, Migration{
			Version:  version,
			Name:     name,
			SQL:      string(body),
			Checksum: hex.EncodeToString(sum[:]),
		})
	}
	slices.SortFunc(out, func(a, b Migration) int { return a.Version - b.Version })
	return out, nil
}

// parseFilename splits "0001_server_settings.sql" into 1 and
// "server_settings".
func parseFilename(filename string) (int, string, error) {
	base := strings.TrimSuffix(filename, ".sql")
	prefix, name, ok := strings.Cut(base, "_")
	if !ok {
		return 0, "", fmt.Errorf("migration %s: expected NNNN_description.sql", filename)
	}
	version, err := strconv.Atoi(prefix)
	if err != nil {
		return 0, "", fmt.Errorf("migration %s: %q is not a version number", filename, prefix)
	}
	if version <= 0 {
		return 0, "", fmt.Errorf("migration %s: version must be positive", filename)
	}
	return version, name, nil
}

// Status lists the migrations already applied.
func Status(ctx context.Context, db *sql.DB, d Dialect) ([]Applied, error) {
	if _, err := db.ExecContext(ctx, d.CreateTable); err != nil {
		return nil, fmt.Errorf("create schema_migrations: %w", err)
	}
	rows, err := db.QueryContext(ctx, d.SelectApplied)
	if err != nil {
		return nil, fmt.Errorf("read schema_migrations: %w", err)
	}
	defer rows.Close()

	var out []Applied
	for rows.Next() {
		var a Applied
		if err := rows.Scan(&a.Version, &a.Name, &a.Checksum, &a.AppliedAt); err != nil {
			return nil, fmt.Errorf("scan schema_migrations: %w", err)
		}
		out = append(out, a)
	}
	return out, rows.Err()
}

// Run applies every migration that has not run yet and returns how many it
// applied. It is safe to call on every boot.
func Run(ctx context.Context, db *sql.DB, d Dialect, ms []Migration, log *slog.Logger) (int, error) {
	if log == nil {
		log = slog.Default()
	}
	applied, err := Status(ctx, db, d)
	if err != nil {
		return 0, err
	}
	known := make(map[int]Applied, len(applied))
	for _, a := range applied {
		known[a.Version] = a
	}

	count := 0
	for _, m := range ms {
		if prev, ok := known[m.Version]; ok {
			if prev.Checksum != m.Checksum {
				return count, fmt.Errorf(
					"migration %04d_%s has changed since it was applied (recorded %s, file %s); "+
						"write a new migration instead of editing an applied one",
					m.Version, m.Name, short(prev.Checksum), short(m.Checksum))
			}
			continue
		}
		if err := apply(ctx, db, d, m); err != nil {
			return count, err
		}
		log.Info("applied migration", "version", m.Version, "name", m.Name)
		count++
	}
	return count, nil
}

func apply(ctx context.Context, db *sql.DB, d Dialect, m Migration) error {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin migration %04d_%s: %w", m.Version, m.Name, err)
	}
	defer func() { _ = tx.Rollback() }()

	if _, err := tx.ExecContext(ctx, m.SQL); err != nil {
		return fmt.Errorf("migration %04d_%s: %w", m.Version, m.Name, err)
	}
	if _, err := tx.ExecContext(ctx, d.InsertApplied, m.Version, m.Name, m.Checksum, time.Now().UTC()); err != nil {
		return fmt.Errorf("record migration %04d_%s: %w", m.Version, m.Name, err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit migration %04d_%s: %w", m.Version, m.Name, err)
	}
	return nil
}

// Pending reports the migrations in ms that have not been applied yet.
func Pending(ctx context.Context, db *sql.DB, d Dialect, ms []Migration) ([]Migration, error) {
	applied, err := Status(ctx, db, d)
	if err != nil {
		return nil, err
	}
	done := make(map[int]struct{}, len(applied))
	for _, a := range applied {
		done[a.Version] = struct{}{}
	}
	var out []Migration
	for _, m := range ms {
		if _, ok := done[m.Version]; !ok {
			out = append(out, m)
		}
	}
	return out, nil
}

// ErrDirtySchema is returned when the recorded schema is ahead of the
// binary: an older warp must not run against a newer database.
var ErrDirtySchema = errors.New("database schema is newer than this binary")

// CheckNotAhead fails when the database records a migration this binary does
// not know about.
func CheckNotAhead(applied []Applied, ms []Migration) error {
	highest := 0
	for _, m := range ms {
		highest = max(highest, m.Version)
	}
	for _, a := range applied {
		if a.Version > highest {
			return fmt.Errorf("%w: database is at %04d_%s, this binary knows up to %04d",
				ErrDirtySchema, a.Version, a.Name, highest)
		}
	}
	return nil
}

func short(checksum string) string {
	if len(checksum) > 8 {
		return checksum[:8]
	}
	return checksum
}