silo-server/internal/database/migrate.go

package database

import (
	"context"
	"fmt"
	"io/fs"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

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

// MigrationFile represents a parsed migration file with its version number,
// filename, and SQL content.
type MigrationFile struct {
	Version  int
	Filename string
	SQL      string
}

const schemaMigrationsLockID int64 = 8_034_219_741

type migrationDB interface {
	Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
	BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error)
}

// RunMigrations applies all pending .up.sql migrations from the given embed.FS.
// It creates a schema_versions table to track which migrations have been applied,
// parses migration filenames to extract version numbers, and applies each
// pending migration within its own transaction.
//
// The dir parameter specifies the subdirectory within the fs.FS that contains
// the migration files.
func RunMigrations(ctx context.Context, pool *pgxpool.Pool, fsys fs.FS, dir string) (err error) {
	conn, err := pool.Acquire(ctx)
	if err != nil {
		return fmt.Errorf("acquiring migration connection: %w", err)
	}
	defer conn.Release()

	if err := lockMigrations(ctx, conn); err != nil {
		return fmt.Errorf("acquiring migration lock: %w", err)
	}
	defer func() {
		if unlockErr := unlockMigrations(conn); unlockErr != nil {
			if err != nil {
				err = fmt.Errorf("%w; releasing migration lock: %v", err, unlockErr)
				return
			}
			err = fmt.Errorf("releasing migration lock: %w", unlockErr)
		}
	}()

	if err := ensureSchemaVersionsTable(ctx, conn); err != nil {
		return fmt.Errorf("ensuring schema_versions table: %w", err)
	}

	migrations, err := loadMigrations(fsys, dir)
	if err != nil {
		return fmt.Errorf("loading migrations: %w", err)
	}

	applied, err := getAppliedVersions(ctx, conn)
	if err != nil {
		return fmt.Errorf("getting applied versions: %w", err)
	}

	for _, m := range migrations {
		if applied[m.Version] {
			continue
		}

		if err := applyMigration(ctx, conn, m); err != nil {
			return fmt.Errorf("applying migration %s (version %d): %w", m.Filename, m.Version, err)
		}
		applied[m.Version] = true
	}

	return nil
}

func lockMigrations(ctx context.Context, db migrationDB) error {
	_, err := db.Exec(ctx, "SELECT pg_advisory_lock($1)", schemaMigrationsLockID)
	return err
}

func unlockMigrations(db migrationDB) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := db.Exec(ctx, "SELECT pg_advisory_unlock($1)", schemaMigrationsLockID)
	return err
}

// ensureSchemaVersionsTable creates the schema_versions table if it does not
// already exist.
func ensureSchemaVersionsTable(ctx context.Context, db migrationDB) error {
	_, err := db.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS public.schema_versions (
			version    INT PRIMARY KEY,
			filename   TEXT NOT NULL,
			applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
		)
	`)
	return err
}

// loadMigrations reads all .up.sql files from the specified directory within
// the filesystem and returns them sorted by version number.
func loadMigrations(fsys fs.FS, dir string) ([]MigrationFile, error) {
	entries, err := fs.ReadDir(fsys, dir)
	if err != nil {
		return nil, fmt.Errorf("reading migrations directory %q: %w", dir, err)
	}

	var migrations []MigrationFile
	seenVersions := make(map[int]string)

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		name := entry.Name()
		if !strings.HasSuffix(name, ".up.sql") {
			continue
		}

		version, err := parseVersion(name)
		if err != nil {
			return nil, fmt.Errorf("parsing version from %q: %w", name, err)
		}
		if existing, ok := seenVersions[version]; ok {
			return nil, fmt.Errorf("duplicate migration version %d: %s and %s", version, existing, name)
		}
		seenVersions[version] = name

		content, err := fs.ReadFile(fsys, filepath.Join(dir, name))
		if err != nil {
			return nil, fmt.Errorf("reading migration file %q: %w", name, err)
		}

		migrations = append(migrations, MigrationFile{
			Version:  version,
			Filename: name,
			SQL:      string(content),
		})
	}

	sort.Slice(migrations, func(i, j int) bool {
		return migrations[i].Version < migrations[j].Version
	})

	return migrations, nil
}

// parseVersion extracts the numeric version prefix from a migration filename.
// For example, "001_initial_schema.up.sql" returns 1.
func parseVersion(filename string) (int, error) {
	// Split on the first underscore to get the version prefix.
	parts := strings.SplitN(filename, "_", 2)
	if len(parts) < 2 {
		return 0, fmt.Errorf("migration filename %q must be in format NNN_description.up.sql", filename)
	}

	version, err := strconv.Atoi(parts[0])
	if err != nil {
		return 0, fmt.Errorf("invalid version number in %q: %w", filename, err)
	}

	if version <= 0 {
		return 0, fmt.Errorf("version number must be positive, got %d in %q", version, filename)
	}

	return version, nil
}

// getAppliedVersions returns a set of migration versions that have already
// been applied.
func getAppliedVersions(ctx context.Context, db migrationDB) (map[int]bool, error) {
	rows, err := db.Query(ctx, "SELECT version FROM public.schema_versions")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	applied := make(map[int]bool)
	for rows.Next() {
		var version int
		if err := rows.Scan(&version); err != nil {
			return nil, err
		}
		applied[version] = true
	}

	return applied, rows.Err()
}

// applyMigration runs a single migration within a transaction and records the
// version in schema_versions.
func applyMigration(ctx context.Context, db migrationDB, m MigrationFile) error {
	tx, err := db.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return fmt.Errorf("beginning transaction: %w", err)
	}
	defer func() {
		// Rollback is a no-op if the tx has already been committed.
		_ = tx.Rollback(ctx)
	}()

	if _, err := tx.Exec(ctx, m.SQL); err != nil {
		return fmt.Errorf("executing SQL: %w", err)
	}

	if _, err := tx.Exec(ctx,
		"INSERT INTO public.schema_versions (version, filename) VALUES ($1, $2)",
		m.Version, m.Filename,
	); err != nil {
		return fmt.Errorf("recording version: %w", err)
	}

	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("committing transaction: %w", err)
	}

	return nil
}