package migration import ( "context" "database/sql" "errors" "fmt" "math" "sync" "time" ) var ( ErrChecksumMismatch = errors.New("migration checksum mismatch") ErrUnknownAppliedMigration = errors.New("unknown applied migration") ErrLockUnavailable = errors.New("migration lock unavailable") sqliteProcessLock sync.Mutex ) const migrationLockID int64 = 0x534d41494c type Applied struct { Version int64 Name string Checksum string AppliedAt time.Time ExecutionMS int64 } type Report struct { CurrentVersion int64 Applied []Applied } type Runner struct { DB *sql.DB Dialect Dialect Migrations []Migration LockWait time.Duration Now func() time.Time } type executor interface { ExecContext(context.Context, string, ...any) (sql.Result, error) QueryContext(context.Context, string, ...any) (*sql.Rows, error) QueryRowContext(context.Context, string, ...any) *sql.Row } func (r Runner) Run(ctx context.Context) (Report, error) { if r.DB == nil { return Report{}, errors.New("migration database is nil") } if _, err := ParseDialect(string(r.Dialect)); err != nil { return Report{}, err } migrations, err := normalize(r.Migrations) if err != nil { return Report{}, err } if err = validateDialectMigrations(migrations, r.Dialect); err != nil { return Report{}, err } conn, err := r.DB.Conn(ctx) if err != nil { return Report{}, fmt.Errorf("open migration connection: %w", err) } defer conn.Close() if r.Dialect == DialectSQLite { return r.runSQLite(ctx, conn, migrations) } if err = r.acquireAdvisoryLock(ctx, conn); err != nil { return Report{}, err } defer r.releaseAdvisoryLock(conn) return r.runTransactional(ctx, conn, migrations) } func validateDialectMigrations(migrations []Migration, dialect Dialect) error { for _, item := range migrations { if item.Apply == nil && len(item.StatementsFor(dialect)) == 0 { return fmt.Errorf("migration %d has no statements for %s", item.Version, dialect) } } return nil } func (r Runner) CurrentVersion(ctx context.Context) (int64, error) { if r.DB == nil { return 0, errors.New("migration database is nil") } ddl, err := schemaTableDDL(r.Dialect) if err != nil { return 0, err } if _, err = r.DB.ExecContext(ctx, ddl); err != nil { return 0, fmt.Errorf("create migration table: %w", err) } var version sql.NullInt64 if err = r.DB.QueryRowContext(ctx, "SELECT MAX(version) FROM schema_migration").Scan(&version); err != nil { return 0, fmt.Errorf("query migration version: %w", err) } if !version.Valid { return 0, nil } return version.Int64, nil } func (r Runner) runSQLite(ctx context.Context, conn *sql.Conn, migrations []Migration) (report Report, resultErr error) { sqliteProcessLock.Lock() defer sqliteProcessLock.Unlock() if _, err := conn.ExecContext(ctx, "BEGIN IMMEDIATE"); err != nil { return Report{}, fmt.Errorf("acquire SQLite migration lock: %w", err) } defer func() { if resultErr != nil { _, _ = conn.ExecContext(context.Background(), "ROLLBACK") } }() report, resultErr = r.applyWithExecutor(ctx, conn, migrations) if resultErr != nil { return Report{}, resultErr } if _, resultErr = conn.ExecContext(ctx, "COMMIT"); resultErr != nil { return Report{}, fmt.Errorf("commit SQLite migrations: %w", resultErr) } return report, nil } func (r Runner) runTransactional(ctx context.Context, conn *sql.Conn, migrations []Migration) (Report, error) { ddl, err := schemaTableDDL(r.Dialect) if err != nil { return Report{}, err } if _, err = conn.ExecContext(ctx, ddl); err != nil { return Report{}, fmt.Errorf("create migration table: %w", err) } applied, err := loadApplied(ctx, conn) if err != nil { return Report{}, err } if err = verifyApplied(applied, migrations); err != nil { return Report{}, err } report := Report{CurrentVersion: highestVersion(applied)} for _, item := range migrations { if _, ok := applied[item.Version]; ok { continue } tx, beginErr := conn.BeginTx(ctx, nil) if beginErr != nil { return Report{}, fmt.Errorf("begin migration %d: %w", item.Version, beginErr) } entry, applyErr := r.applyOne(ctx, tx, item) if applyErr != nil { _ = tx.Rollback() return Report{}, applyErr } if commitErr := tx.Commit(); commitErr != nil { return Report{}, fmt.Errorf("commit migration %d: %w", item.Version, commitErr) } report.Applied = append(report.Applied, entry) report.CurrentVersion = item.Version } return report, nil } func (r Runner) applyWithExecutor(ctx context.Context, exec executor, migrations []Migration) (Report, error) { ddl, err := schemaTableDDL(r.Dialect) if err != nil { return Report{}, err } if _, err = exec.ExecContext(ctx, ddl); err != nil { return Report{}, fmt.Errorf("create migration table: %w", err) } applied, err := loadApplied(ctx, exec) if err != nil { return Report{}, err } if err = verifyApplied(applied, migrations); err != nil { return Report{}, err } report := Report{CurrentVersion: highestVersion(applied)} for _, item := range migrations { if _, ok := applied[item.Version]; ok { continue } entry, applyErr := r.applyOne(ctx, exec, item) if applyErr != nil { return Report{}, applyErr } report.Applied = append(report.Applied, entry) report.CurrentVersion = item.Version } return report, nil } func (r Runner) applyOne(ctx context.Context, exec executor, item Migration) (Applied, error) { started := time.Now() if r.Now != nil { started = r.Now() } if item.Apply != nil { if err := item.Apply(ctx, exec, r.Dialect); err != nil { return Applied{}, fmt.Errorf("execute migration %d (%s): %w", item.Version, item.Name, err) } } else { for _, statement := range item.StatementsFor(r.Dialect) { if _, err := exec.ExecContext(ctx, statement); err != nil { return Applied{}, fmt.Errorf("execute migration %d (%s): %w", item.Version, item.Name, err) } } } finished := time.Now() if r.Now != nil { finished = r.Now() } executionMS := int64(math.Max(0, float64(finished.Sub(started).Milliseconds()))) entry := Applied{Version: item.Version, Name: item.Name, Checksum: item.Checksum(), AppliedAt: finished.UTC(), ExecutionMS: executionMS} query, err := insertSQL(r.Dialect) if err != nil { return Applied{}, err } if _, err = exec.ExecContext(ctx, query, entry.Version, entry.Name, entry.Checksum, entry.AppliedAt, entry.ExecutionMS); err != nil { return Applied{}, fmt.Errorf("record migration %d: %w", item.Version, err) } return entry, nil } func loadApplied(ctx context.Context, exec executor) (map[int64]Applied, error) { rows, err := exec.QueryContext(ctx, "SELECT version, name, checksum, applied_at, execution_ms FROM schema_migration ORDER BY version") if err != nil { return nil, fmt.Errorf("query applied migrations: %w", err) } defer rows.Close() result := map[int64]Applied{} for rows.Next() { var item Applied if err = rows.Scan(&item.Version, &item.Name, &item.Checksum, &item.AppliedAt, &item.ExecutionMS); err != nil { return nil, fmt.Errorf("scan applied migration: %w", err) } result[item.Version] = item } if err = rows.Err(); err != nil { return nil, fmt.Errorf("iterate applied migrations: %w", err) } return result, nil } func verifyApplied(applied map[int64]Applied, migrations []Migration) error { known := make(map[int64]Migration, len(migrations)) for _, item := range migrations { known[item.Version] = item } for version, stored := range applied { item, ok := known[version] if !ok { return fmt.Errorf("%w: version %d", ErrUnknownAppliedMigration, version) } if stored.Name != item.Name || !item.AcceptsChecksum(stored.Checksum) { return fmt.Errorf("%w: version %d", ErrChecksumMismatch, version) } } return nil } func highestVersion(applied map[int64]Applied) int64 { var result int64 for version := range applied { if version > result { result = version } } return result } func (r Runner) acquireAdvisoryLock(ctx context.Context, conn *sql.Conn) error { wait := r.LockWait if wait <= 0 { wait = 30 * time.Second } switch r.Dialect { case DialectMySQL: var acquired sql.NullInt64 if err := conn.QueryRowContext(ctx, "SELECT GET_LOCK(?, ?)", "smail_schema_migration", int(math.Ceil(wait.Seconds()))).Scan(&acquired); err != nil { return fmt.Errorf("acquire MySQL migration lock: %w", err) } if !acquired.Valid || acquired.Int64 != 1 { return ErrLockUnavailable } case DialectPostgres: deadline := time.Now().Add(wait) for { var acquired bool if err := conn.QueryRowContext(ctx, "SELECT pg_try_advisory_lock($1)", migrationLockID).Scan(&acquired); err != nil { return fmt.Errorf("acquire PostgreSQL migration lock: %w", err) } if acquired { return nil } remaining := time.Until(deadline) if remaining <= 0 { return ErrLockUnavailable } delay := 100 * time.Millisecond if remaining < delay { delay = remaining } timer := time.NewTimer(delay) select { case <-ctx.Done(): if !timer.Stop() { <-timer.C } return ctx.Err() case <-timer.C: } } default: return fmt.Errorf("unsupported advisory lock dialect %q", r.Dialect) } return nil } func (r Runner) releaseAdvisoryLock(conn *sql.Conn) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() switch r.Dialect { case DialectMySQL: _, _ = conn.ExecContext(ctx, "SELECT RELEASE_LOCK(?)", "smail_schema_migration") case DialectPostgres: _, _ = conn.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", migrationLockID) } }