feat: 重构启动迁移逻辑

Author: quantal

backend/internal/bootstrap/startup.go

@@ -2,17 +2,8 @@ package bootstrap
 
 import (
 	"context"
-	"crypto/sha256"
-	stdsql "database/sql"
-	"embed"
-	"encoding/hex"
-	"fmt"
-	"io/fs"
 	"log/slog"
-	"path/filepath"
-	"sort"
 	"strings"
-	"time"
 
 	entsql "entgo.io/ent/dialect/sql"
 
@@ -23,277 +14,17 @@ import (
 	"github.com/DouDOU-start/airgate-core/internal/auth"
 )
 
-//go:embed migrations/*.sql
-var startupMigrationFiles embed.FS
-
 // RunStartupTasks 运行启动阶段的整理任务。
 func RunStartupTasks(db *ent.Client, drv *entsql.Driver, apiKeySecret string) {
 	slog.Info("bootstrap_startup_tasks_start")
 	backfillKeyHints(db, apiKeySecret)
 	backfillResellerMarkupColumns(drv)
 	migrateAccountState(drv)
 	migrateUserHistoryRefs(drv)
-	runStartupMigrations(drv)
+	runSystemUpgrades(drv)
 	slog.Info("bootstrap_startup_tasks_done")
 }
 
-type startupMigration struct {
-	ID          string
-	Description string
-	Checksum    string
-	SQL         string
-}
-
-func runStartupMigrations(drv *entsql.Driver) {
-	if drv == nil {
-		return
-	}
-	migrations := loadStartupMigrations()
-	if len(migrations) == 0 {
-		return
-	}
-
-	ctx := context.Background()
-	conn, err := drv.DB().Conn(ctx)
-	if err != nil {
-		panicStartupMigration("open startup migration connection", err)
-	}
-	defer func() { _ = conn.Close() }()
-
-	const createTableSQL = `CREATE TABLE IF NOT EXISTS public.bootstrap_migrations (
-		id text PRIMARY KEY,
-		description text NOT NULL DEFAULT '',
-		checksum text NOT NULL DEFAULT '',
-		applied_at timestamptz NOT NULL DEFAULT now(),
-		duration_ms bigint NOT NULL DEFAULT 0
-	)`
-	if _, err := conn.ExecContext(ctx, createTableSQL); err != nil {
-		panicStartupMigration("create bootstrap_migrations table", err)
-	}
-	if _, err := conn.ExecContext(ctx, `ALTER TABLE public.bootstrap_migrations ADD COLUMN IF NOT EXISTS checksum text NOT NULL DEFAULT ''`); err != nil {
-		panicStartupMigration("add bootstrap_migrations checksum column", err)
-	}
-
-	const lockKey int64 = 2026052809517
-	if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1)`, lockKey); err != nil {
-		panicStartupMigration("lock startup migrations", err)
-	}
-	defer func() {
-		if _, err := conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock($1)`, lockKey); err != nil {
-			slog.Warn("bootstrap_migration_unlock_failed", sdk.LogFieldError, err)
-		}
-	}()
-
-	for _, migration := range migrations {
-		var appliedChecksum stdsql.NullString
-		const appliedSQL = `SELECT checksum FROM public.bootstrap_migrations WHERE id = $1`
-		err := conn.QueryRowContext(ctx, appliedSQL, migration.ID).Scan(&appliedChecksum)
-		if err == nil {
-			if appliedChecksum.Valid && appliedChecksum.String != "" && appliedChecksum.String != migration.Checksum {
-				panicStartupMigration("verify startup migration checksum "+migration.ID, fmt.Errorf("recorded=%s current=%s", appliedChecksum.String, migration.Checksum))
-			}
-			if !appliedChecksum.Valid || appliedChecksum.String == "" {
-				const updateSQL = `UPDATE public.bootstrap_migrations
-					SET checksum = $2, description = $3
-					WHERE id = $1 AND checksum = ''`
-				if _, err := conn.ExecContext(ctx, updateSQL, migration.ID, migration.Checksum, migration.Description); err != nil {
-					panicStartupMigration("backfill startup migration checksum "+migration.ID, err)
-				}
-			}
-			continue
-		}
-		if err != stdsql.ErrNoRows {
-			panicStartupMigration("check startup migration "+migration.ID, err)
-		}
-
-		start := time.Now()
-		slog.Info("bootstrap_migration_start", "id", migration.ID)
-		if err := executeStartupMigrationSQL(ctx, conn, migration); err != nil {
-			panicStartupMigration("run startup migration "+migration.ID, err)
-		}
-		duration := time.Since(start).Milliseconds()
-		const insertSQL = `INSERT INTO public.bootstrap_migrations (id, description, checksum, duration_ms)
-			VALUES ($1, $2, $3, $4)
-			ON CONFLICT (id) DO NOTHING`
-		if _, err := conn.ExecContext(ctx, insertSQL, migration.ID, migration.Description, migration.Checksum, duration); err != nil {
-			panicStartupMigration("record startup migration "+migration.ID, err)
-		}
-		slog.Info("bootstrap_migration_done", "id", migration.ID, "duration_ms", duration)
-	}
-}
-
-func loadStartupMigrations() []startupMigration {
-	entries, err := fs.ReadDir(startupMigrationFiles, "migrations")
-	if err != nil {
-		panicStartupMigration("read startup migration files", err)
-	}
-	sort.Slice(entries, func(i, j int) bool {
-		return entries[i].Name() < entries[j].Name()
-	})
-
-	migrations := make([]startupMigration, 0, len(entries))
-	for _, entry := range entries {
-		if entry.IsDir() || filepath.Ext(entry.Name()) != ".sql" {
-			continue
-		}
-		path := filepath.ToSlash(filepath.Join("migrations", entry.Name()))
-		data, err := startupMigrationFiles.ReadFile(path)
-		if err != nil {
-			panicStartupMigration("read startup migration "+entry.Name(), err)
-		}
-		hash := sha256.Sum256(data)
-		id := strings.TrimSuffix(entry.Name(), ".sql")
-		sql := string(data)
-		migrations = append(migrations, startupMigration{
-			ID:          id,
-			Description: startupMigrationDescription(sql, id),
-			Checksum:    hex.EncodeToString(hash[:]),
-			SQL:         sql,
-		})
-	}
-	return migrations
-}
-
-func startupMigrationDescription(sql, fallback string) string {
-	for _, line := range strings.Split(sql, "\n") {
-		line = strings.TrimSpace(line)
-		if strings.HasPrefix(line, "-- description:") {
-			return strings.TrimSpace(strings.TrimPrefix(line, "-- description:"))
-		}
-	}
-	return fallback
-}
-
-func executeStartupMigrationSQL(ctx context.Context, conn *stdsql.Conn, migration startupMigration) error {
-	for _, stmt := range splitSQLStatements(migration.SQL) {
-		if _, err := conn.ExecContext(ctx, stmt); err != nil {
-			return fmt.Errorf("execute statement in %s: %w", migration.ID, err)
-		}
-	}
-	return nil
-}
-
-func splitSQLStatements(sql string) []string {
-	var statements []string
-	start := 0
-	inSingleQuote := false
-	inDoubleQuote := false
-	inLineComment := false
-	inBlockComment := false
-	dollarTag := ""
-
-	for i := 0; i < len(sql); i++ {
-		ch := sql[i]
-		next := byte(0)
-		if i+1 < len(sql) {
-			next = sql[i+1]
-		}
-
-		switch {
-		case inLineComment:
-			if ch == '\n' {
-				inLineComment = false
-			}
-			continue
-		case inBlockComment:
-			if ch == '*' && next == '/' {
-				inBlockComment = false
-				i++
-			}
-			continue
-		case dollarTag != "":
-			if strings.HasPrefix(sql[i:], dollarTag) {
-				i += len(dollarTag) - 1
-				dollarTag = ""
-			}
-			continue
-		case inSingleQuote:
-			if ch == '\'' {
-				if next == '\'' {
-					i++
-				} else {
-					inSingleQuote = false
-				}
-			}
-			continue
-		case inDoubleQuote:
-			if ch == '"' {
-				if next == '"' {
-					i++
-				} else {
-					inDoubleQuote = false
-				}
-			}
-			continue
-		}
-
-		if ch == '-' && next == '-' {
-			inLineComment = true
-			i++
-			continue
-		}
-		if ch == '/' && next == '*' {
-			inBlockComment = true
-			i++
-			continue
-		}
-		if ch == '\'' {
-			inSingleQuote = true
-			continue
-		}
-		if ch == '"' {
-			inDoubleQuote = true
-			continue
-		}
-		if ch == '$' {
-			if tag, ok := sqlDollarTag(sql[i:]); ok {
-				dollarTag = tag
-				i += len(tag) - 1
-				continue
-			}
-		}
-		if ch == ';' {
-			stmt := strings.TrimSpace(sql[start:i])
-			if stmt != "" {
-				statements = append(statements, stmt)
-			}
-			start = i + 1
-		}
-	}
-
-	tail := strings.TrimSpace(sql[start:])
-	if tail != "" {
-		statements = append(statements, tail)
-	}
-	return statements
-}
-
-func sqlDollarTag(sql string) (string, bool) {
-	if sql == "" || sql[0] != '$' {
-		return "", false
-	}
-	for i := 1; i < len(sql); i++ {
-		ch := sql[i]
-		if ch == '$' {
-			return sql[:i+1], true
-		}
-		if !((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '_') {
-			return "", false
-		}
-	}
-	return "", false
-}
-
-func panicStartupMigration(action string, err error) {
-	if err != nil {
-		err = fmt.Errorf("%s: %w", action, err)
-	} else {
-		err = fmt.Errorf("%s", action)
-	}
-	slog.Error("bootstrap_startup_migration_failed", sdk.LogFieldError, err)
-	panic(err)
-}
-
 // migrateUserHistoryRefs 允许硬删除用户，同时保留历史使用记录和余额流水。
 // 用量/计费聚合依赖 usage_logs 的成本快照字段；这里把历史表的 user 外键改为 SET NULL，
 // 并回填 user_id/user_email 快照，避免删除用户后历史记录丢失归属信息。