uPortal2/backend/internal/migration/migrator.go

package migration

import (
	"context"
	"database/sql"
	"log/slog"
	"sort"
	"time"
)

type Migration struct {
	Version int64
	SQL     string
}

type packageMigration struct {
	pkg       string
	migration Migration
}

type Migrator struct {
	db   *sql.DB
	list []packageMigration
}

func NewMigrator(db *sql.DB) *Migrator {
	return &Migrator{db: db}
}

func (m *Migrator) Register(pkg string, migrations ...Migration) {
	for _, migration := range migrations {
		m.list = append(m.list, packageMigration{pkg: pkg, migration: migration})
	}
}

func (m *Migrator) Run(ctx context.Context) error {
	if _, err := m.db.ExecContext(ctx, "PRAGMA foreign_keys = ON"); err != nil {
		slog.Info("enable foreign keys failed", "error", err)
		return ErrEnableForeignKeys
	}

	const createTracking = `CREATE TABLE IF NOT EXISTS schema_migrations (
    package    TEXT    NOT NULL,
    version    INTEGER NOT NULL,
    applied_at TEXT    NOT NULL,
    PRIMARY KEY (package, version)
);`
	if _, err := m.db.ExecContext(ctx, createTracking); err != nil {
		slog.Info("create schema_migrations failed", "error", err)
		return ErrCreateSchemaMigrations
	}

	m.sort()

	for _, entry := range m.list {
		ok, err := m.isApplied(ctx, entry)
		if err != nil {
			slog.Info("check migration failed", "package", entry.pkg, "version", entry.migration.Version, "error", err)
			return ErrCheckMigration
		}
		if ok {
			continue
		}
		if err := m.apply(ctx, entry); err != nil {
			return err
		}
	}

	return nil
}

func (m *Migrator) sort() {
	order := make(map[string]int)
	for _, entry := range m.list {
		if _, ok := order[entry.pkg]; !ok {
			order[entry.pkg] = len(order)
		}
	}
	sort.SliceStable(m.list, func(i, j int) bool {
		if m.list[i].pkg != m.list[j].pkg {
			return order[m.list[i].pkg] < order[m.list[j].pkg]
		}
		return m.list[i].migration.Version < m.list[j].migration.Version
	})
}

func (m *Migrator) isApplied(ctx context.Context, entry packageMigration) (bool, error) {
	const query = "SELECT 1 FROM schema_migrations WHERE package = ? AND version = ?"
	var one int
	err := m.db.QueryRowContext(ctx, query, entry.pkg, entry.migration.Version).Scan(&one)
	if err == sql.ErrNoRows {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}

func (m *Migrator) apply(ctx context.Context, entry packageMigration) error {
	fail := func(err error) error {
		slog.Info("apply migration failed", "package", entry.pkg, "version", entry.migration.Version, "error", err)
		return ErrMigrationFailed
	}

	tx, err := m.db.BeginTx(ctx, nil)
	if err != nil {
		return fail(err)
	}

	if _, err := tx.ExecContext(ctx, entry.migration.SQL); err != nil {
		_ = tx.Rollback()
		return fail(err)
	}

	const insertTracking = "INSERT INTO schema_migrations(package, version, applied_at) VALUES(?, ?, ?)"
	if _, err := tx.ExecContext(ctx, insertTracking, entry.pkg, entry.migration.Version, time.Now().UTC()); err != nil {
		_ = tx.Rollback()
		return fail(err)
	}

	if err := tx.Commit(); err != nil {
		return fail(err)
	}

	return nil
}