package migration import ( "context" "database/sql" "log/slog" "sort" "time" ) type Migration struct { Version int64 SQL string } type packageMigration struct { pkg string migration Migration } type Migrator struct { db *sql.DB list []packageMigration } func NewMigrator(db *sql.DB) *Migrator { return &Migrator{db: db} } func (m *Migrator) Register(pkg string, migrations ...Migration) { for _, migration := range migrations { m.list = append(m.list, packageMigration{pkg: pkg, migration: migration}) } } func (m *Migrator) Run(ctx context.Context) error { if _, err := m.db.ExecContext(ctx, "PRAGMA foreign_keys = ON"); err != nil { slog.Info("enable foreign keys failed", "error", err) return ErrEnableForeignKeys } const createTracking = `CREATE TABLE IF NOT EXISTS schema_migrations ( package TEXT NOT NULL, version INTEGER NOT NULL, applied_at TEXT NOT NULL, PRIMARY KEY (package, version) );` if _, err := m.db.ExecContext(ctx, createTracking); err != nil { slog.Info("create schema_migrations failed", "error", err) return ErrCreateSchemaMigrations } m.sort() for _, entry := range m.list { ok, err := m.isApplied(ctx, entry) if err != nil { slog.Info("check migration failed", "package", entry.pkg, "version", entry.migration.Version, "error", err) return ErrCheckMigration } if ok { continue } if err := m.apply(ctx, entry); err != nil { return err } } return nil } func (m *Migrator) sort() { order := make(map[string]int) for _, entry := range m.list { if _, ok := order[entry.pkg]; !ok { order[entry.pkg] = len(order) } } sort.SliceStable(m.list, func(i, j int) bool { if m.list[i].pkg != m.list[j].pkg { return order[m.list[i].pkg] < order[m.list[j].pkg] } return m.list[i].migration.Version < m.list[j].migration.Version }) } func (m *Migrator) isApplied(ctx context.Context, entry packageMigration) (bool, error) { const query = "SELECT 1 FROM schema_migrations WHERE package = ? AND version = ?" var one int err := m.db.QueryRowContext(ctx, query, entry.pkg, entry.migration.Version).Scan(&one) if err == sql.ErrNoRows { return false, nil } if err != nil { return false, err } return true, nil } func (m *Migrator) apply(ctx context.Context, entry packageMigration) error { fail := func(err error) error { slog.Info("apply migration failed", "package", entry.pkg, "version", entry.migration.Version, "error", err) return ErrMigrationFailed } tx, err := m.db.BeginTx(ctx, nil) if err != nil { return fail(err) } if _, err := tx.ExecContext(ctx, entry.migration.SQL); err != nil { _ = tx.Rollback() return fail(err) } const insertTracking = "INSERT INTO schema_migrations(package, version, applied_at) VALUES(?, ?, ?)" if _, err := tx.ExecContext(ctx, insertTracking, entry.pkg, entry.migration.Version, time.Now().UTC()); err != nil { _ = tx.Rollback() return fail(err) } if err := tx.Commit(); err != nil { return fail(err) } return nil }