From 0b4f7bf56588c1060375f094406fe530b7086dcf Mon Sep 17 00:00:00 2001
From: Ryan Olds <ryan.olds@oden.io>
Date: Wed, 15 Jun 2022 14:33:35 -0700
Subject: [PATCH] feat(migrate): added option to only mark migration up/down as
 applied on success

Fixed bug with returning wrong err var

Changed the field name to be more clear

fix: Mark applied/unapplied before up/down when marking applied on success and failure

fix: bug with inverted condition
---
 chmigrate/migrator.go | 41 ++++++++++++++++++++++++++++++++---------
 1 file changed, 32 insertions(+), 9 deletions(-)

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
 			}
 		}
 	}