feat(migrate): add WithMarkAppliedOnSuccess

migrate/migrator.go

@@ -26,14 +26,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(enabled bool) MigratorOption {
+	return func(m *Migrator) {
+		m.markAppliedOnSuccess = enabled
+	}
+}
+
 type Migrator struct {
 	db         *bun.DB
 	migrations *Migrations
 
 	ms MigrationSlice
 
-	table      string
-	locksTable string
+	table                string
+	locksTable           string
+	markAppliedOnSuccess bool
 }
 
 func NewMigrator(db *bun.DB, migrations *Migrations, opts ...MigratorOption) *Migrator {
@@ -148,9 +157,10 @@ func (m *Migrator) Migrate(ctx context.Context, opts ...MigrationOption) (*Migra
 		migration := &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 group, err
+		if !m.markAppliedOnSuccess {
+			if err := m.MarkApplied(ctx, migration); err != nil {
+				return group, err
+			}
 		}
 
 		group.Migrations = migrations[:i+1]
@@ -160,6 +170,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
@@ -187,14 +203,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
 			}
 		}
 	}