knowledge-governance/internal/platform/postgres/migrate.go

package postgres

import (
	"context"
	"embed"
	"fmt"
	"io/fs"
	"sort"
	"strconv"
	"strings"

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

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

type migration struct {
	Version int64
	UpSQL   string
	DownSQL string
}

func loadMigrations() ([]migration, error) {
	paths, err := fs.Glob(migrationFS, "migrations/*.up.sql")
	if err != nil {
		return nil, fmt.Errorf("glob migrations: %w", err)
	}
	result := make([]migration, 0, len(paths))
	for _, path := range paths {
		base := strings.TrimPrefix(path, "migrations/")
		versionText, _, ok := strings.Cut(base, "_")
		if !ok {
			return nil, fmt.Errorf("invalid migration filename %q", path)
		}
		version, err := strconv.ParseInt(versionText, 10, 64)
		if err != nil {
			return nil, fmt.Errorf("parse migration version %q: %w", path, err)
		}
		up, err := migrationFS.ReadFile(path)
		if err != nil {
			return nil, fmt.Errorf("read up migration %q: %w", path, err)
		}
		downPath := strings.TrimSuffix(path, ".up.sql") + ".down.sql"
		down, err := migrationFS.ReadFile(downPath)
		if err != nil {
			return nil, fmt.Errorf("read down migration %q: %w", downPath, err)
		}
		result = append(result, migration{Version: version, UpSQL: string(up), DownSQL: string(down)})
	}
	sort.Slice(result, func(i, j int) bool { return result[i].Version < result[j].Version })
	return result, nil
}

const migrationLockID int64 = 7346729101

func ApplyMigrations(ctx context.Context, pool *pgxpool.Pool) error {
	if _, err := pool.Exec(ctx, "SELECT pg_advisory_lock($1)", migrationLockID); err != nil {
		return fmt.Errorf("lock migrations: %w", err)
	}
	defer func() { _, _ = pool.Exec(context.Background(), "SELECT pg_advisory_unlock($1)", migrationLockID) }()

	if _, err := pool.Exec(ctx, `CREATE TABLE IF NOT EXISTS schema_migrations (
		version BIGINT PRIMARY KEY,
		applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
	)`); err != nil {
		return fmt.Errorf("create migration ledger: %w", err)
	}

	migrations, err := loadMigrations()
	if err != nil {
		return err
	}
	for _, item := range migrations {
		var applied bool
		if err := pool.QueryRow(ctx, "SELECT EXISTS (SELECT 1 FROM schema_migrations WHERE version = $1)", item.Version).Scan(&applied); err != nil {
			return fmt.Errorf("check migration %d: %w", item.Version, err)
		}
		if applied {
			continue
		}
		tx, err := pool.BeginTx(ctx, pgx.TxOptions{})
		if err != nil {
			return fmt.Errorf("begin migration %d: %w", item.Version, err)
		}
		if _, err := tx.Exec(ctx, item.UpSQL); err != nil {
			_ = tx.Rollback(ctx)
			return fmt.Errorf("apply migration %d: %w", item.Version, err)
		}
		if _, err := tx.Exec(ctx, "INSERT INTO schema_migrations (version) VALUES ($1)", item.Version); err != nil {
			_ = tx.Rollback(ctx)
			return fmt.Errorf("record migration %d: %w", item.Version, err)
		}
		if err := tx.Commit(ctx); err != nil {
			return fmt.Errorf("commit migration %d: %w", item.Version, err)
		}
	}
	return nil
}

func RollbackLatestMigration(ctx context.Context, pool *pgxpool.Pool) error {
	migrations, err := loadMigrations()
	if err != nil {
		return err
	}
	var version int64
	if err := pool.QueryRow(ctx, "SELECT version FROM schema_migrations ORDER BY version DESC LIMIT 1").Scan(&version); err != nil {
		if err == pgx.ErrNoRows {
			return nil
		}
		return fmt.Errorf("read latest migration: %w", err)
	}
	var selected *migration
	for i := range migrations {
		if migrations[i].Version == version {
			selected = &migrations[i]
			break
		}
	}
	if selected == nil {
		return fmt.Errorf("missing down migration for version %d", version)
	}
	tx, err := pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return fmt.Errorf("begin rollback %d: %w", version, err)
	}
	if _, err := tx.Exec(ctx, selected.DownSQL); err != nil {
		_ = tx.Rollback(ctx)
		return fmt.Errorf("rollback migration %d: %w", version, err)
	}
	if _, err := tx.Exec(ctx, "DELETE FROM schema_migrations WHERE version = $1", version); err != nil {
		_ = tx.Rollback(ctx)
		return fmt.Errorf("delete migration ledger %d: %w", version, err)
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit rollback %d: %w", version, err)
	}
	return nil
}