ai-gateway-go/internal/platform/migrate/migrate.go

package migrate

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

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

const advisoryLockID int64 = 6720240805

type Migration struct {
	Version  string
	Filename string
	SQL      string
	Checksum string
}

func Load(directory string) ([]Migration, error) {
	entries, err := os.ReadDir(directory)
	if err != nil {
		return nil, fmt.Errorf("read migration directory: %w", err)
	}
	var filenames []string
	for _, entry := range entries {
		if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".sql") {
			filenames = append(filenames, entry.Name())
		}
	}
	sort.Strings(filenames)
	migrations := make([]Migration, 0, len(filenames))
	for _, filename := range filenames {
		body, err := os.ReadFile(filepath.Join(directory, filename))
		if err != nil {
			return nil, fmt.Errorf("read migration %s: %w", filename, err)
		}
		version, _, ok := strings.Cut(filename, "_")
		if !ok || version == "" {
			return nil, fmt.Errorf("migration %s must start with a version and underscore", filename)
		}
		digest := sha256.Sum256(body)
		migrations = append(migrations, Migration{
			Version: version, Filename: filename, SQL: string(body), Checksum: hex.EncodeToString(digest[:]),
		})
	}
	return migrations, nil
}

func Apply(ctx context.Context, pool *pgxpool.Pool, migrations []Migration) error {
	connection, err := pool.Acquire(ctx)
	if err != nil {
		return fmt.Errorf("acquire migration connection: %w", err)
	}
	defer connection.Release()
	if _, err := connection.Exec(ctx, "SELECT pg_advisory_lock($1)", advisoryLockID); err != nil {
		return fmt.Errorf("lock migrations: %w", err)
	}
	defer func() { _, _ = connection.Exec(context.Background(), "SELECT pg_advisory_unlock($1)", advisoryLockID) }()

	if _, err := connection.Exec(ctx, `
        CREATE SCHEMA IF NOT EXISTS gateway;
        CREATE TABLE IF NOT EXISTS gateway.schema_migrations (
            version text PRIMARY KEY,
            filename text NOT NULL,
            checksum text NOT NULL,
            applied_at timestamptz NOT NULL DEFAULT clock_timestamp()
        )`); err != nil {
		return fmt.Errorf("initialize migration table: %w", err)
	}

	for _, migration := range migrations {
		var existingChecksum string
		err := connection.QueryRow(ctx,
			"SELECT checksum FROM gateway.schema_migrations WHERE version = $1", migration.Version,
		).Scan(&existingChecksum)
		if err == nil {
			if existingChecksum != migration.Checksum {
				return fmt.Errorf("migration %s checksum changed after application", migration.Filename)
			}
			continue
		}
		if !errors.Is(err, pgx.ErrNoRows) {
			return fmt.Errorf("check migration %s: %w", migration.Filename, err)
		}
		transaction, err := connection.Begin(ctx)
		if err != nil {
			return fmt.Errorf("begin migration %s: %w", migration.Filename, err)
		}
		if _, err := transaction.Exec(ctx, migration.SQL); err != nil {
			_ = transaction.Rollback(ctx)
			return fmt.Errorf("apply migration %s: %w", migration.Filename, err)
		}
		if _, err := transaction.Exec(ctx,
			"INSERT INTO gateway.schema_migrations (version, filename, checksum) VALUES ($1, $2, $3)",
			migration.Version, migration.Filename, migration.Checksum,
		); err != nil {
			_ = transaction.Rollback(ctx)
			return fmt.Errorf("record migration %s: %w", migration.Filename, err)
		}
		if err := transaction.Commit(ctx); err != nil {
			return fmt.Errorf("commit migration %s: %w", migration.Filename, err)
		}
	}
	return nil
}