From 34a35450ffe80127812aa231cc17b5318f018366 Mon Sep 17 00:00:00 2001
From: Marcel Gregoriadis <m.gregoriadis@gmail.com>
Date: Mon, 5 Sep 2022 14:03:54 +0300
Subject: [PATCH] Add --number flag to dbmate down

---
 main.go          |  9 ++++++-
 pkg/dbmate/db.go | 67 +++++++++++++++++++++++-------------------------
 2 files changed, 40 insertions(+), 36 deletions(-)

diff --git a/main.go b/main.go
index 04a777f5..1988ceff 100644
--- a/main.go
+++ b/main.go
@@ -150,10 +150,17 @@ func NewApp() *cli.App {
 					EnvVars: []string{"DBMATE_VERBOSE"},
 					Usage:   "print the result of each statement execution",
 				},
+				&cli.IntFlag{
+					Name:    "number",
+					Aliases: []string{"n"},
+					Usage:   "number of migrations to rollback (-1 means all)",
+					Value:   1,
+				},
 			},
 			Action: action(func(db *dbmate.DB, c *cli.Context) error {
 				db.Verbose = c.Bool("verbose")
-				return db.Rollback()
+				number := c.Int("number")
+				return db.Rollback(number)
 			}),
 		},
 		{
diff --git a/pkg/dbmate/db.go b/pkg/dbmate/db.go
index 5c1c3a28..7728adc6 100644
--- a/pkg/dbmate/db.go
+++ b/pkg/dbmate/db.go
@@ -464,7 +464,7 @@ func migrationVersion(filename string) string {
 }
 
 // Rollback rolls back the most recent migration
-func (db *DB) Rollback() error {
+func (db *DB) Rollback(number int) error {
 	drv, err := db.GetDriver()
 	if err != nil {
 		return err
@@ -483,55 +483,52 @@ func (db *DB) Rollback() error {
 	}
 	defer dbutil.MustClose(sqlDB)
 
-	applied, err := drv.SelectMigrations(sqlDB, 1)
+	applied, err := drv.SelectMigrations(sqlDB, number)
 	if err != nil {
 		return err
 	}
 
-	// grab most recent applied migration (applied has len=1)
-	latest := ""
-	for ver := range applied {
-		latest = ver
-	}
-	if latest == "" {
+	if len(applied) == 0 {
 		return ErrNoRollback
 	}
 
-	filename, err := findMigrationFile(db.MigrationsDir, latest)
-	if err != nil {
-		return err
-	}
-
-	fmt.Fprintf(db.Log, "Rolling back: %s\n", filename)
+	for ver := range applied {
+		filename, err := findMigrationFile(db.MigrationsDir, ver)
+		if err != nil {
+			return err
+		}
 
-	_, down, err := parseMigration(filepath.Join(db.MigrationsDir, filename))
-	if err != nil {
-		return err
-	}
+		fmt.Fprintf(db.Log, "Rolling back: %s\n", filename)
 
-	execMigration := func(tx dbutil.Transaction) error {
-		// rollback migration
-		result, err := tx.Exec(down.Contents)
+		_, down, err := parseMigration(filepath.Join(db.MigrationsDir, filename))
 		if err != nil {
 			return err
-		} else if db.Verbose {
-			db.printVerbose(result)
 		}
 
-		// remove migration record
-		return drv.DeleteMigration(tx, latest)
-	}
+		execMigration := func(tx dbutil.Transaction) error {
+			// rollback migration
+			result, err := tx.Exec(down.Contents)
+			if err != nil {
+				return err
+			} else if db.Verbose {
+				db.printVerbose(result)
+			}
 
-	if down.Options.Transaction() {
-		// begin transaction
-		err = doTransaction(sqlDB, execMigration)
-	} else {
-		// run outside of transaction
-		err = execMigration(sqlDB)
-	}
+			// remove migration record
+			return drv.DeleteMigration(tx, ver)
+		}
 
-	if err != nil {
-		return err
+		if down.Options.Transaction() {
+			// begin transaction
+			err = doTransaction(sqlDB, execMigration)
+		} else {
+			// run outside of transaction
+			err = execMigration(sqlDB)
+		}
+
+		if err != nil {
+			return err
+		}
 	}
 
 	// automatically update schema file, silence errors