2 files changed, 53 insertions, 20 deletions

diff --git a/vendor/github.com/uptrace/bun/migrate/migration.go b/vendor/github.com/uptrace/bun/migrate/migration.go
index 05bd6006c..ae649446d 100644
--- a/vendor/github.com/uptrace/bun/migrate/migration.go
+++ b/vendor/github.com/uptrace/bun/migrate/migration.go
@@ -25,11 +25,11 @@ type Migration struct {
 	Down MigrationFunc `bun:"-"`
 }
 
-func (m *Migration) String() string {
+func (m Migration) String() string {
 	return m.Name
 }
 
-func (m *Migration) IsApplied() bool {
+func (m Migration) IsApplied() bool {
 	return m.ID > 0
 }
 
@@ -89,6 +89,22 @@ func NewSQLMigrationFunc(fsys fs.FS, name string) MigrationFunc {
 			idb = conn
 		}
 
+		var retErr error
+
+		defer func() {
+			if tx, ok := idb.(bun.Tx); ok {
+				retErr = tx.Commit()
+				return
+			}
+
+			if conn, ok := idb.(bun.Conn); ok {
+				retErr = conn.Close()
+				return
+			}
+
+			panic("not reached")
+		}()
+
 		for _, q := range queries {
 			_, err = idb.ExecContext(ctx, q)
 			if err != nil {
@@ -96,13 +112,7 @@ func NewSQLMigrationFunc(fsys fs.FS, name string) MigrationFunc {
 			}
 		}
 
-		if tx, ok := idb.(bun.Tx); ok {
-			return tx.Commit()
-		} else if conn, ok := idb.(bun.Conn); ok {
-			return conn.Close()
-		}
-
-		panic("not reached")
+		return retErr
 	}
 }
 
@@ -222,11 +232,11 @@ type MigrationGroup struct {
 	Migrations MigrationSlice
 }
 
-func (g *MigrationGroup) IsZero() bool {
+func (g MigrationGroup) IsZero() bool {
 	return g.ID == 0 && len(g.Migrations) == 0
 }
 
-func (g *MigrationGroup) String() string {
+func (g MigrationGroup) String() string {
 	if g.IsZero() {
 		return "nil"
 	}
diff --git a/vendor/github.com/uptrace/bun/migrate/migrator.go b/vendor/github.com/uptrace/bun/migrate/migrator.go
index e271b7a06..d0efca2fc 100644
--- a/vendor/github.com/uptrace/bun/migrate/migrator.go
+++ b/vendor/github.com/uptrace/bun/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
 			}
 		}
 	}