axe/internal/db/migrate.go

package db

import (
	"context"
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"fmt"
	"io/fs"
	"log/slog"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"

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

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

const migrationsDir = "migrations"

// Имя файла: <version>_<name>.sql, например 0001_init.sql.
var migrationFileRe = regexp.MustCompile(`^(\d+)_([a-z0-9_-]+)\.sql$`)

type migration struct {
	Version  int
	Name     string
	Checksum string
	SQL      string
}

// Migrate применяет недостающие миграции в порядке версии.
// Каждая миграция выполняется в отдельной транзакции; применённые
// фиксируются в schema_migrations вместе с checksum файла.
func Migrate(ctx context.Context, pool *pgxpool.Pool) error {
	migrations, err := loadMigrations()
	if err != nil {
		return err
	}

	if _, err := pool.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS schema_migrations (
			version    int PRIMARY KEY,
			name       text NOT NULL,
			checksum   text NOT NULL,
			applied_at timestamptz NOT NULL DEFAULT now()
		)`); err != nil {
		return fmt.Errorf("create schema_migrations: %w", err)
	}

	applied, err := appliedVersions(ctx, pool)
	if err != nil {
		return err
	}

	for _, m := range migrations {
		checksum, ok := applied[m.Version]
		if !ok {
			if err := applyMigration(ctx, pool, m); err != nil {
				return fmt.Errorf("migration %04d_%s: %w", m.Version, m.Name, err)
			}
			slog.Info("migration applied", "version", m.Version, "name", m.Name)
			continue
		}
		if checksum != m.Checksum {
			// Лучше упасть на старте, чем работать с непредсказуемой схемой.
			return fmt.Errorf(
				"migration %04d_%s: checksum changed after apply (db=%s, file=%s); "+
					"applied migrations must not be edited — add a new migration instead",
				m.Version, m.Name, checksum[:12], m.Checksum[:12],
			)
		}
	}
	return nil
}

// AppliedVersions — версии, которые уже накатаны. Экспортируется для тестов.
func AppliedVersions(ctx context.Context, pool *pgxpool.Pool) (map[int]string, error) {
	return appliedVersions(ctx, pool)
}

func appliedVersions(ctx context.Context, pool *pgxpool.Pool) (map[int]string, error) {
	rows, err := pool.Query(ctx, `SELECT version, checksum FROM schema_migrations`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := map[int]string{}
	for rows.Next() {
		var version int
		var checksum string
		if err := rows.Scan(&version, &checksum); err != nil {
			return nil, err
		}
		out[version] = checksum
	}
	return out, rows.Err()
}

func applyMigration(ctx context.Context, pool *pgxpool.Pool, m migration) error {
	tx, err := pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx) //nolint:errcheck // откат не нужен после Commit

	// Без аргументов pgx уходит в simple protocol — разрешает несколько
	// statement'ов в одной строке, что удобно для DDL-файлов.
	if _, err := tx.Exec(ctx, m.SQL); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx,
		`INSERT INTO schema_migrations (version, name, checksum) VALUES ($1, $2, $3)`,
		m.Version, m.Name, m.Checksum,
	); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func loadMigrations() ([]migration, error) {
	entries, err := fs.ReadDir(migrationsFS, migrationsDir)
	if err != nil {
		return nil, fmt.Errorf("read migrations: %w", err)
	}

	var out []migration
	seen := map[int]string{}

	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sql") {
			continue
		}

		match := migrationFileRe.FindStringSubmatch(entry.Name())
		if match == nil {
			return nil, fmt.Errorf("migration %q: filename must be <version>_<name>.sql", entry.Name())
		}

		version, err := strconv.Atoi(match[1])
		if err != nil {
			return nil, fmt.Errorf("migration %q: bad version: %w", entry.Name(), err)
		}
		if prev, ok := seen[version]; ok {
			return nil, fmt.Errorf("migration version %d used twice: %s and %s", version, prev, entry.Name())
		}
		seen[version] = entry.Name()

		body, err := fs.ReadFile(migrationsFS, path.Join(migrationsDir, entry.Name()))
		if err != nil {
			return nil, fmt.Errorf("migration %q: %w", entry.Name(), err)
		}
		sum := sha256.Sum256(body)

		out = append(out, migration{
			Version:  version,
			Name:     match[2],
			Checksum: hex.EncodeToString(sum[:]),
			SQL:      string(body),
		})
	}

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