hyapp-server/server/admin/internal/migration/migration.go

package migration

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

func Apply(ctx context.Context, db *sql.DB, dir string) error {
	if strings.TrimSpace(dir) == "" {
		return errors.New("migration dir is required")
	}
	if _, err := os.Stat(dir); err != nil {
		return err
	}
	if err := ensureSchemaTable(ctx, db); err != nil {
		return err
	}

	files, err := migrationFiles(dir)
	if err != nil {
		return err
	}
	for _, file := range files {
		version := filepath.Base(file)
		applied, err := isApplied(ctx, db, version, file)
		if err != nil {
			return err
		}
		if applied {
			continue
		}
		if err := applyFile(ctx, db, version, file); err != nil {
			return err
		}
	}
	return nil
}

func ensureSchemaTable(ctx context.Context, db *sql.DB) error {
	if _, err := db.ExecContext(ctx, `
CREATE TABLE IF NOT EXISTS schema_migrations (
  version VARCHAR(128) PRIMARY KEY,
  checksum CHAR(64) NOT NULL DEFAULT '',
  dirty BOOLEAN NOT NULL DEFAULT FALSE,
  applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci`); err != nil {
		return err
	}
	if err := ensureColumn(ctx, db, "checksum", "ALTER TABLE schema_migrations ADD COLUMN checksum CHAR(64) NOT NULL DEFAULT '' AFTER version"); err != nil {
		return err
	}
	return ensureColumn(ctx, db, "dirty", "ALTER TABLE schema_migrations ADD COLUMN dirty BOOLEAN NOT NULL DEFAULT FALSE AFTER checksum")
}

func ensureColumn(ctx context.Context, db *sql.DB, column string, alter string) error {
	var count int
	err := db.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM information_schema.columns
WHERE table_schema = DATABASE() AND table_name = 'schema_migrations' AND column_name = ?`, column).Scan(&count)
	if err != nil {
		return err
	}
	if count > 0 {
		return nil
	}
	_, err = db.ExecContext(ctx, alter)
	return err
}

func migrationFiles(dir string) ([]string, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	files := make([]string, 0, len(entries))
	for _, entry := range entries {
		if entry.IsDir() || filepath.Ext(entry.Name()) != ".sql" {
			continue
		}
		files = append(files, filepath.Join(dir, entry.Name()))
	}
	sort.Strings(files)
	return files, nil
}

func isApplied(ctx context.Context, db *sql.DB, version string, file string) (bool, error) {
	var row migrationRow
	err := db.QueryRowContext(ctx, "SELECT version, checksum, dirty FROM schema_migrations WHERE version = ?", version).Scan(&row.Version, &row.Checksum, &row.Dirty)
	if errors.Is(err, sql.ErrNoRows) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	if row.Dirty {
		return false, fmt.Errorf("migration %s is dirty; repair it before continuing", version)
	}
	body, err := os.ReadFile(file)
	if err != nil {
		return false, err
	}
	sum := checksum(body)
	if row.Checksum != "" && row.Checksum != sum {
		return false, fmt.Errorf("migration %s checksum changed", version)
	}
	if row.Checksum == "" {
		if _, err := db.ExecContext(ctx, "UPDATE schema_migrations SET checksum = ? WHERE version = ?", sum, version); err != nil {
			return false, err
		}
	}
	return true, nil
}

func applyFile(ctx context.Context, db *sql.DB, version string, file string) error {
	body, err := os.ReadFile(file)
	if err != nil {
		return err
	}
	sum := checksum(body)

	var existing migrationRow
	err = db.QueryRowContext(ctx, "SELECT version, checksum, dirty FROM schema_migrations WHERE version = ?", version).Scan(&existing.Version, &existing.Checksum, &existing.Dirty)
	if err == nil {
		if existing.Dirty {
			return fmt.Errorf("migration %s is dirty; repair it before continuing", version)
		}
		if existing.Checksum != "" && existing.Checksum != sum {
			return fmt.Errorf("migration %s checksum changed", version)
		}
		return nil
	}
	if !errors.Is(err, sql.ErrNoRows) {
		return err
	}

	if _, err := db.ExecContext(ctx, "INSERT INTO schema_migrations(version, checksum, dirty) VALUES (?, ?, TRUE)", version, sum); err != nil {
		return err
	}

	for _, statement := range splitSQL(string(body)) {
		if _, err := db.ExecContext(ctx, statement); err != nil {
			return fmt.Errorf("apply migration %s: %w", version, err)
		}
	}
	if _, err := db.ExecContext(ctx, "UPDATE schema_migrations SET dirty = FALSE, checksum = ?, applied_at = CURRENT_TIMESTAMP WHERE version = ?", sum, version); err != nil {
		return err
	}
	return nil
}

func splitSQL(raw string) []string {
	statements := make([]string, 0)
	var b strings.Builder
	var quote rune
	var inLineComment bool
	var inBlockComment bool
	escaped := false
	for i, r := range raw {
		next := rune(0)
		if i+1 < len(raw) {
			next = rune(raw[i+1])
		}
		if inLineComment {
			b.WriteRune(r)
			if r == '\n' {
				inLineComment = false
			}
			continue
		}
		if inBlockComment {
			b.WriteRune(r)
			if r == '*' && next == '/' {
				inBlockComment = false
			}
			continue
		}
		if quote == 0 && r == '-' && next == '-' {
			inLineComment = true
			b.WriteRune(r)
			continue
		}
		if quote == 0 && r == '/' && next == '*' {
			inBlockComment = true
			b.WriteRune(r)
			continue
		}
		if quote != 0 {
			b.WriteRune(r)
			if escaped {
				escaped = false
				continue
			}
			if r == '\\' {
				escaped = true
				continue
			}
			if r == quote {
				quote = 0
			}
			continue
		}
		switch r {
		case '\'', '"', '`':
			quote = r
			b.WriteRune(r)
		case ';':
			appendStatement(&statements, b.String())
			b.Reset()
		default:
			b.WriteRune(r)
		}
	}
	appendStatement(&statements, b.String())
	return statements
}

type migrationRow struct {
	Version  string
	Checksum string
	Dirty    bool
}

func checksum(body []byte) string {
	sum := sha256.Sum256(body)
	return hex.EncodeToString(sum[:])
}

func appendStatement(statements *[]string, raw string) {
	statement := strings.TrimSpace(raw)
	if statement == "" {
		return
	}
	*statements = append(*statements, statement)
}