feat(db): SQLite 连接与迁移框架（embed 迁移/事务校验/PG 预留 M4）

2026-09-06 04:26:02 +08:00
1 parent dd16176bb7
commit 34cac241c1
@@ -0,0 +1,219 @@
+package db
+
+import (
+	"context"
+	"database/sql"
+	"embed"
+	"fmt"
+	"log/slog"
+	"sort"
+	"strings"
+)
+
+//go:embed migrations/*.sql
+var migrationsFS embed.FS
+
+// migration is a single schema change with its forward (up) and reverse (down)
+// SQL. id is the "{timestamp}_{desc}" derived from the file name, which is also
+// the unique key recorded in schema_migrations.
+type migration struct {
+	id      string
+	upSQL   string
+	downSQL string
+}
+
+// createMigrationsTable is a parameter-free DDL statement (no placeholders). It
+// is executed with ExecContext to create the bookkeeping table on first use.
+const createMigrationsTable = `
+CREATE TABLE IF NOT EXISTS schema_migrations (
+    id         TEXT PRIMARY KEY,
+    applied_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
+)`
+
+// MigrateUp applies every embedded migration that has not yet been recorded in
+// schema_migrations. Migrations run in ascending timestamp order, each inside its
+// own transaction; the schema_migrations row is inserted in that same
+// transaction, so a failing migration leaves no half-applied state. Already
+// applied migrations are skipped, making MigrateUp idempotent.
+//
+// The entire content of a migration's .up.sql is executed as one ExecContext
+// inside the transaction. Keep migrations small and use several files for a
+// larger schema change rather than packing many statements into one file.
+func MigrateUp(ctx context.Context, db *sql.DB, log *slog.Logger) error {
+	if err := ensureMigrationTable(ctx, db); err != nil {
+		return err
+	}
+	applied, err := appliedMigrationIDs(ctx, db)
+	if err != nil {
+		return err
+	}
+	migs, err := loadMigrations()
+	if err != nil {
+		return err
+	}
+
+	for _, m := range migs {
+		if applied[m.id] {
+			continue
+		}
+		if err := applyUp(ctx, db, m, discardLogger(log)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// MigrateDown rolls back the most recent steps migrations, in descending
+// timestamp order. Each rollback runs the migration's .down.sql and deletes its
+// schema_migrations row in the same transaction. steps must be >= 1; if steps
+// exceeds the number of applied migrations, all applied migrations are rolled
+// back.
+func MigrateDown(ctx context.Context, db *sql.DB, steps int, log *slog.Logger) error {
+	if steps < 1 {
+		return fmt.Errorf("db: down: %w", ErrInvalidSteps)
+	}
+	log = discardLogger(log)
+
+	if err := ensureMigrationTable(ctx, db); err != nil {
+		return err
+	}
+	applied, err := appliedMigrationIDs(ctx, db)
+	if err != nil {
+		return err
+	}
+	migs, err := loadMigrations()
+	if err != nil {
+		return err
+	}
+
+	toRollback := make([]migration, 0, len(applied))
+	for _, m := range migs {
+		if applied[m.id] {
+			toRollback = append(toRollback, m)
+		}
+	}
+	// Reverse to descending order, then keep only the most recent `steps`.
+	for i, j := 0, len(toRollback)-1; i < j; i, j = i+1, j-1 {
+		toRollback[i], toRollback[j] = toRollback[j], toRollback[i]
+	}
+	if len(toRollback) > steps {
+		toRollback = toRollback[:steps]
+	}
+
+	for _, m := range toRollback {
+		if err := applyDown(ctx, db, m, log); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// ensureMigrationTable creates the schema_migrations bookkeeping table if it does
+// not already exist.
+func ensureMigrationTable(ctx context.Context, db *sql.DB) error {
+	if _, err := db.ExecContext(ctx, createMigrationsTable); err != nil {
+		return fmt.Errorf("db: ensure schema_migrations: %w", err)
+	}
+	return nil
+}
+
+// appliedMigrationIDs returns the set of migration ids already recorded in
+// schema_migrations.
+func appliedMigrationIDs(ctx context.Context, db *sql.DB) (map[string]bool, error) {
+	rows, err := db.QueryContext(ctx, "SELECT id FROM schema_migrations")
+	if err != nil {
+		return nil, fmt.Errorf("db: list applied migrations: %w", err)
+	}
+	defer rows.Close()
+
+	applied := make(map[string]bool)
+	for rows.Next() {
+		var id string
+		if err := rows.Scan(&id); err != nil {
+			return nil, fmt.Errorf("db: scan migration id: %w", err)
+		}
+		applied[id] = true
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("db: iterate applied migrations: %w", err)
+	}
+	return applied, nil
+}
+
+// applyUp runs a single migration's up SQL in a transaction and records it.
+func applyUp(ctx context.Context, db *sql.DB, m migration, log *slog.Logger) error {
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		return fmt.Errorf("db: begin up %s: %w", m.id, err)
+	}
+	defer tx.Rollback()
+
+	if _, err := tx.ExecContext(ctx, m.upSQL); err != nil {
+		return fmt.Errorf("db: apply up %s: %w", m.id, err)
+	}
+	if _, err := tx.ExecContext(ctx, "INSERT INTO schema_migrations (id) VALUES (?)", m.id); err != nil {
+		return fmt.Errorf("db: record up %s: %w", m.id, err)
+	}
+	if err := tx.Commit(); err != nil {
+		return fmt.Errorf("db: commit up %s: %w", m.id, err)
+	}
+
+	log.Info("migration applied", "migration", m.id)
+	return nil
+}
+
+// applyDown runs a single migration's down SQL in a transaction and unrecords it.
+func applyDown(ctx context.Context, db *sql.DB, m migration, log *slog.Logger) error {
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		return fmt.Errorf("db: begin down %s: %w", m.id, err)
+	}
+	defer tx.Rollback()
+
+	if _, err := tx.ExecContext(ctx, m.downSQL); err != nil {
+		return fmt.Errorf("db: apply down %s: %w", m.id, err)
+	}
+	if _, err := tx.ExecContext(ctx, "DELETE FROM schema_migrations WHERE id = ?", m.id); err != nil {
+		return fmt.Errorf("db: unrecord down %s: %w", m.id, err)
+	}
+	if err := tx.Commit(); err != nil {
+		return fmt.Errorf("db: commit down %s: %w", m.id, err)
+	}
+
+	log.Info("migration rolled back", "migration", m.id)
+	return nil
+}
+
+// loadMigrations reads all embedded .up.sql files, pairs each with its matching
+// .down.sql, and returns them sorted ascending by id (which places them in
+// timestamp order). A .up.sql without a matching .down.sql is an error.
+func loadMigrations() ([]migration, error) {
+	entries, err := migrationsFS.ReadDir("migrations")
+	if err != nil {
+		return nil, fmt.Errorf("db: read embedded migrations: %w", err)
+	}
+
+	migs := make([]migration, 0, len(entries)/2)
+	for _, e := range entries {
+		if e.IsDir() || !strings.HasSuffix(e.Name(), ".up.sql") {
+			continue
+		}
+		id := strings.TrimSuffix(e.Name(), ".up.sql")
+		upSQL, err := migrationsFS.ReadFile("migrations/" + e.Name())
+		if err != nil {
+			return nil, fmt.Errorf("db: read %s: %w", e.Name(), err)
+		}
+		downName := id + ".down.sql"
+		downSQL, err := migrationsFS.ReadFile("migrations/" + downName)
+		if err != nil {
+			return nil, fmt.Errorf("db: %w: %s", ErrDownMissing, downName)
+		}
+		migs = append(migs, migration{id: id, upSQL: string(upSQL), downSQL: string(downSQL)})
+	}
+
+	if len(migs) == 0 {
+		return nil, ErrNoEmbedded
+	}
+	sort.Slice(migs, func(i, j int) bool { return migs[i].id < migs[j].id })
+	return migs, nil
+}