material_manage_system/pkg/migration/runner.go
2026-09-15 16:08:18 +08:00

203 lines
5.0 KiB
Go

package migration
import (
"context"
"database/sql"
"fmt"
"os"
"path/filepath"
"sort"
"strings"
)
const createMigrationsTable = `CREATE TABLE IF NOT EXISTS schema_migrations (
version VARCHAR(100) PRIMARY KEY,
dirty BOOLEAN NOT NULL DEFAULT TRUE,
applied_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
)`
// Run applies .sql files in lexical filename order.
// MySQL 5.7 may implicitly commit DDL, so every migration is marked dirty
// before its first statement and only marked clean after all statements pass.
func Run(ctx context.Context, db *sql.DB, dir string) error {
if db == nil {
return fmt.Errorf("migration database is nil")
}
if _, err := db.ExecContext(ctx, createMigrationsTable); err != nil {
return fmt.Errorf("create schema_migrations: %w", err)
}
var dirtyVersion string
err := db.QueryRowContext(ctx, `SELECT version FROM schema_migrations WHERE dirty = TRUE ORDER BY applied_at DESC LIMIT 1`).Scan(&dirtyVersion)
if err == nil {
return fmt.Errorf("migration %s is dirty; stop and perform manual repair before retrying", dirtyVersion)
}
if err != sql.ErrNoRows {
return fmt.Errorf("check dirty migration: %w", err)
}
applied, err := loadApplied(ctx, db)
if err != nil {
return err
}
files, err := migrationFiles(dir)
if err != nil {
return err
}
for _, path := range files {
version := filepath.Base(path)
if applied[version] {
continue
}
contents, err := os.ReadFile(path)
if err != nil {
return fmt.Errorf("read migration %s: %w", version, err)
}
statements, err := splitStatements(string(contents))
if err != nil {
return fmt.Errorf("parse migration %s: %w", version, err)
}
if _, err := db.ExecContext(ctx, `INSERT INTO schema_migrations (version, dirty) VALUES (?, TRUE)`, version); err != nil {
return fmt.Errorf("mark migration %s dirty: %w", version, err)
}
for index, statement := range statements {
if _, err := db.ExecContext(ctx, statement); err != nil {
return fmt.Errorf("migration %s statement %d failed: %w; version remains dirty because MySQL DDL cannot be transactionally rolled back; perform manual repair", version, index+1, err)
}
}
if _, err := db.ExecContext(ctx, `UPDATE schema_migrations SET dirty = FALSE, applied_at = CURRENT_TIMESTAMP WHERE version = ?`, version); err != nil {
return fmt.Errorf("mark migration %s clean: %w; version remains dirty and requires manual repair", version, err)
}
}
return nil
}
func loadApplied(ctx context.Context, db *sql.DB) (map[string]bool, error) {
rows, err := db.QueryContext(ctx, `SELECT version FROM schema_migrations WHERE dirty = FALSE`)
if err != nil {
return nil, fmt.Errorf("list applied migrations: %w", err)
}
defer rows.Close()
applied := make(map[string]bool)
for rows.Next() {
var version string
if err := rows.Scan(&version); err != nil {
return nil, fmt.Errorf("scan applied migration: %w", err)
}
applied[version] = true
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("iterate applied migrations: %w", err)
}
return applied, nil
}
func migrationFiles(dir string) ([]string, error) {
entries, err := os.ReadDir(dir)
if err != nil {
return nil, fmt.Errorf("read migrations directory %s: %w", dir, err)
}
files := make([]string, 0, len(entries))
for _, entry := range entries {
if entry.IsDir() || strings.ToLower(filepath.Ext(entry.Name())) != ".sql" {
continue
}
files = append(files, filepath.Join(dir, entry.Name()))
}
sort.Strings(files)
return files, nil
}
func splitStatements(input string) ([]string, error) {
var statements []string
var current strings.Builder
var quote rune
escaped := false
lineComment := false
blockComment := false
runes := []rune(input)
flush := func() {
statement := strings.TrimSpace(current.String())
if statement != "" {
statements = append(statements, statement)
}
current.Reset()
}
for i := 0; i < len(runes); i++ {
ch := runes[i]
next := rune(0)
if i+1 < len(runes) {
next = runes[i+1]
}
if lineComment {
if ch == '\n' {
lineComment = false
if current.Len() > 0 {
current.WriteRune(' ')
}
}
continue
}
if blockComment {
if ch == '*' && next == '/' {
blockComment = false
i++
if current.Len() > 0 {
current.WriteRune(' ')
}
}
continue
}
if quote != 0 {
current.WriteRune(ch)
if escaped {
escaped = false
continue
}
if ch == '\\' {
escaped = true
continue
}
if ch == quote {
quote = 0
}
continue
}
if ch == '-' && next == '-' {
lineComment = true
i++
continue
}
if ch == '#' {
lineComment = true
continue
}
if ch == '/' && next == '*' {
blockComment = true
i++
continue
}
if ch == '\'' || ch == '"' || ch == '`' {
quote = ch
current.WriteRune(ch)
continue
}
if ch == ';' {
flush()
continue
}
current.WriteRune(ch)
}
if quote != 0 {
return nil, fmt.Errorf("unterminated quoted string")
}
if blockComment {
return nil, fmt.Errorf("unterminated block comment")
}
flush()
return statements, nil
}