diff --git a/chmigrate/migrator.go b/chmigrate/migrator.go index 2513135..f4a3f98 100644 --- a/chmigrate/migrator.go +++ b/chmigrate/migrator.go @@ -27,14 +27,23 @@ func WithLocksTableName(table string) MigratorOption { } } +// WithMarkAppliedOnSuccess sets the migrator to only mark migrations as applied/unapplied +// when their up/down is successful +func WithMarkAppliedOnSuccess() MigratorOption { + return func(m *Migrator) { + m.markAppliedOnSuccess = true + } +} + type Migrator struct { db *ch.DB migrations *Migrations ms MigrationSlice - table string - locksTable string + table string + locksTable string + markAppliedOnSuccess bool } func NewMigrator(db *ch.DB, migrations *Migrations, opts ...MigratorOption) *Migrator { @@ -149,9 +158,10 @@ func (m *Migrator) Migrate(ctx context.Context, opts ...MigrationOption) (*Migra migration := &group.Migrations[i] migration.GroupID = group.ID - // Always mark migration as applied so the rollback has a chance to fix the database. - if err := m.MarkApplied(ctx, migration); err != nil { - return nil, err + if !m.markAppliedOnSuccess { + if err := m.MarkApplied(ctx, migration); err != nil { + return group, err + } } if !cfg.nop && migration.Up != nil { @@ -159,6 +169,12 @@ func (m *Migrator) Migrate(ctx context.Context, opts ...MigrationOption) (*Migra return group, err } } + + if m.markAppliedOnSuccess { + if err := m.MarkApplied(ctx, migration); err != nil { + return group, err + } + } } return group, nil @@ -186,14 +202,21 @@ func (m *Migrator) Rollback(ctx context.Context, opts ...MigrationOption) (*Migr for i := len(lastGroup.Migrations) - 1; i >= 0; i-- { migration := &lastGroup.Migrations[i] - // Always mark migration as unapplied to match migrate behavior. - if err := m.MarkUnapplied(ctx, migration); err != nil { - return nil, err + if !m.markAppliedOnSuccess { + if err := m.MarkUnapplied(ctx, migration); err != nil { + return lastGroup, err + } } if !cfg.nop && migration.Down != nil { if err := migration.Down(ctx, m.db); err != nil { - return nil, err + return lastGroup, err + } + } + + if m.markAppliedOnSuccess { + if err := m.MarkUnapplied(ctx, migration); err != nil { + return lastGroup, err } } }