Cyrene_v2/internal/db/migrate.go

package db

import (
	"context"
	"database/sql"
	"embed"
	"fmt"
	"log/slog"
	"sort"
	"strings"
)

//go:embed migrations/*.sql
var migrationsFS embed.FS

// migration is a single schema change with its forward (up) and reverse (down)
// SQL. id is the "{timestamp}_{desc}" derived from the file name, which is also
// the unique key recorded in schema_migrations.
type migration struct {
	id      string
	upSQL   string
	downSQL string
}

// createMigrationsTable is a parameter-free DDL statement (no placeholders). It
// is executed with ExecContext to create the bookkeeping table on first use.
const createMigrationsTable = `
CREATE TABLE IF NOT EXISTS schema_migrations (
    id         TEXT PRIMARY KEY,
    applied_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
)`

// MigrateUp applies every embedded migration that has not yet been recorded in
// schema_migrations. Migrations run in ascending timestamp order, each inside its
// own transaction; the schema_migrations row is inserted in that same
// transaction, so a failing migration leaves no half-applied state. Already
// applied migrations are skipped, making MigrateUp idempotent.
//
// The entire content of a migration's .up.sql is executed as one ExecContext
// inside the transaction. Keep migrations small and use several files for a
// larger schema change rather than packing many statements into one file.
func MigrateUp(ctx context.Context, db *sql.DB, log *slog.Logger) error {
	if err := ensureMigrationTable(ctx, db); err != nil {
		return err
	}
	applied, err := appliedMigrationIDs(ctx, db)
	if err != nil {
		return err
	}
	migs, err := loadMigrations()
	if err != nil {
		return err
	}

	for _, m := range migs {
		if applied[m.id] {
			continue
		}
		if err := applyUp(ctx, db, m, discardLogger(log)); err != nil {
			return err
		}
	}
	return nil
}

// MigrateDown rolls back the most recent steps migrations, in descending
// timestamp order. Each rollback runs the migration's .down.sql and deletes its
// schema_migrations row in the same transaction. steps must be >= 1; if steps
// exceeds the number of applied migrations, all applied migrations are rolled
// back.
func MigrateDown(ctx context.Context, db *sql.DB, steps int, log *slog.Logger) error {
	if steps < 1 {
		return fmt.Errorf("db: down: %w", ErrInvalidSteps)
	}
	log = discardLogger(log)

	if err := ensureMigrationTable(ctx, db); err != nil {
		return err
	}
	applied, err := appliedMigrationIDs(ctx, db)
	if err != nil {
		return err
	}
	migs, err := loadMigrations()
	if err != nil {
		return err
	}

	toRollback := make([]migration, 0, len(applied))
	for _, m := range migs {
		if applied[m.id] {
			toRollback = append(toRollback, m)
		}
	}
	// Reverse to descending order, then keep only the most recent `steps`.
	for i, j := 0, len(toRollback)-1; i < j; i, j = i+1, j-1 {
		toRollback[i], toRollback[j] = toRollback[j], toRollback[i]
	}
	if len(toRollback) > steps {
		toRollback = toRollback[:steps]
	}

	for _, m := range toRollback {
		if err := applyDown(ctx, db, m, log); err != nil {
			return err
		}
	}
	return nil
}

// ensureMigrationTable creates the schema_migrations bookkeeping table if it does
// not already exist.
func ensureMigrationTable(ctx context.Context, db *sql.DB) error {
	if _, err := db.ExecContext(ctx, createMigrationsTable); err != nil {
		return fmt.Errorf("db: ensure schema_migrations: %w", err)
	}
	return nil
}

// appliedMigrationIDs returns the set of migration ids already recorded in
// schema_migrations.
func appliedMigrationIDs(ctx context.Context, db *sql.DB) (map[string]bool, error) {
	rows, err := db.QueryContext(ctx, "SELECT id FROM schema_migrations")
	if err != nil {
		return nil, fmt.Errorf("db: list applied migrations: %w", err)
	}
	defer rows.Close()

	applied := make(map[string]bool)
	for rows.Next() {
		var id string
		if err := rows.Scan(&id); err != nil {
			return nil, fmt.Errorf("db: scan migration id: %w", err)
		}
		applied[id] = true
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("db: iterate applied migrations: %w", err)
	}
	return applied, nil
}

// applyUp runs a single migration's up SQL in a transaction and records it.
func applyUp(ctx context.Context, db *sql.DB, m migration, log *slog.Logger) error {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("db: begin up %s: %w", m.id, err)
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx, m.upSQL); err != nil {
		return fmt.Errorf("db: apply up %s: %w", m.id, err)
	}
	if _, err := tx.ExecContext(ctx, "INSERT INTO schema_migrations (id) VALUES (?)", m.id); err != nil {
		return fmt.Errorf("db: record up %s: %w", m.id, err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("db: commit up %s: %w", m.id, err)
	}

	log.Info("migration applied", "migration", m.id)
	return nil
}

// applyDown runs a single migration's down SQL in a transaction and unrecords it.
func applyDown(ctx context.Context, db *sql.DB, m migration, log *slog.Logger) error {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("db: begin down %s: %w", m.id, err)
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx, m.downSQL); err != nil {
		return fmt.Errorf("db: apply down %s: %w", m.id, err)
	}
	if _, err := tx.ExecContext(ctx, "DELETE FROM schema_migrations WHERE id = ?", m.id); err != nil {
		return fmt.Errorf("db: unrecord down %s: %w", m.id, err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("db: commit down %s: %w", m.id, err)
	}

	log.Info("migration rolled back", "migration", m.id)
	return nil
}

// loadMigrations reads all embedded .up.sql files, pairs each with its matching
// .down.sql, and returns them sorted ascending by id (which places them in
// timestamp order). A .up.sql without a matching .down.sql is an error.
func loadMigrations() ([]migration, error) {
	entries, err := migrationsFS.ReadDir("migrations")
	if err != nil {
		return nil, fmt.Errorf("db: read embedded migrations: %w", err)
	}

	migs := make([]migration, 0, len(entries)/2)
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".up.sql") {
			continue
		}
		id := strings.TrimSuffix(e.Name(), ".up.sql")
		upSQL, err := migrationsFS.ReadFile("migrations/" + e.Name())
		if err != nil {
			return nil, fmt.Errorf("db: read %s: %w", e.Name(), err)
		}
		downName := id + ".down.sql"
		downSQL, err := migrationsFS.ReadFile("migrations/" + downName)
		if err != nil {
			return nil, fmt.Errorf("db: %w: %s", ErrDownMissing, downName)
		}
		migs = append(migs, migration{id: id, upSQL: string(upSQL), downSQL: string(downSQL)})
	}

	if len(migs) == 0 {
		return nil, ErrNoEmbedded
	}
	sort.Slice(migs, func(i, j int) bool { return migs[i].id < migs[j].id })
	return migs, nil
}