266 lines
6.8 KiB
Go
266 lines
6.8 KiB
Go
package migration
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"database/sql"
|
||
"encoding/hex"
|
||
"errors"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"sort"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
// Options controls migration repair behavior.
|
||
type Options struct {
|
||
// AllowChecksumRepair lets local/dev databases accept edited historical migrations.
|
||
// Production-like environments must keep this false so checksum drift fails closed.
|
||
AllowChecksumRepair bool
|
||
}
|
||
|
||
func Apply(ctx context.Context, db *sql.DB, dir string) error {
|
||
return ApplyWithOptions(ctx, db, dir, Options{})
|
||
}
|
||
|
||
func ApplyWithOptions(ctx context.Context, db *sql.DB, dir string, options Options) 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, options)
|
||
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_ms BIGINT NOT NULL DEFAULT 0
|
||
) 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
|
||
}
|
||
if err := ensureColumn(ctx, db, "dirty", "ALTER TABLE schema_migrations ADD COLUMN dirty BOOLEAN NOT NULL DEFAULT FALSE AFTER checksum"); err != nil {
|
||
return err
|
||
}
|
||
return ensureColumn(ctx, db, "applied_at_ms", "ALTER TABLE schema_migrations ADD COLUMN applied_at_ms BIGINT NOT NULL DEFAULT 0 AFTER dirty")
|
||
}
|
||
|
||
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, options Options) (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 {
|
||
if options.AllowChecksumRepair {
|
||
// 本地开发阶段允许直接编辑历史迁移；修复只更新记录，不重放历史 SQL，避免重复 ALTER 破坏现有库。
|
||
if _, err := db.ExecContext(ctx, "UPDATE schema_migrations SET checksum = ? WHERE version = ?", sum, version); err != nil {
|
||
return false, err
|
||
}
|
||
return true, nil
|
||
}
|
||
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_ms = ? WHERE version = ?", sum, time.Now().UTC().UnixMilli(), 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)
|
||
}
|