Add RollbackOnly & MigrateOnly with tests

Author: realkarych

pkg/dbmate/db.go

@@ -741,6 +741,143 @@ func (db *DB) Rollback() error {
 	return nil
 }
 
+// MigrateOnly applies exactly one pending migration that matches the given version.
+func (db *DB) MigrateOnly(version string) error {
+	drv, err := db.Driver()
+	if err != nil {
+		return err
+	}
+
+	migrations, err := db.FindMigrations()
+	if err != nil {
+		return err
+	}
+
+	var target *Migration
+	highestApplied := ""
+	for i, m := range migrations {
+		if m.Version == version {
+			target = &migrations[i]
+		} else if m.Applied && db.Strict && m.Version > highestApplied {
+			highestApplied = m.Version
+		}
+	}
+
+	if target == nil {
+		return ErrMigrationNotFound
+	}
+	if target.Applied {
+		fmt.Fprintf(db.Log, "Skipping (already applied): %s\n", target.FileName)
+		return nil
+	}
+	if db.Strict && target.Version <= highestApplied {
+		return fmt.Errorf("migration `%s` is out of order with applied `%s` in --strict mode", target.Version, highestApplied)
+	}
+
+	sqlDB, err := db.openDatabaseForMigration(drv)
+	if err != nil {
+		return err
+	}
+	defer dbutil.MustClose(sqlDB)
+
+	fmt.Fprintf(db.Log, "Applying: %s\n", target.FileName)
+	start := time.Now()
+
+	parsed, err := target.Parse()
+	if err != nil {
+		return err
+	}
+	exec := func(tx dbutil.Transaction) error {
+		res, err := tx.Exec(parsed.Up)
+		if err != nil {
+			return drv.QueryError(parsed.Up, err)
+		} else if db.Verbose {
+			db.printVerbose(res)
+		}
+		return drv.InsertMigration(tx, target.Version)
+	}
+
+	if parsed.UpOptions.Transaction() {
+		err = doTransaction(sqlDB, exec)
+	} else {
+		err = exec(sqlDB)
+	}
+
+	fmt.Fprintf(db.Log, "Applied: %s in %s\n", target.FileName, time.Since(start))
+	if err != nil {
+		return err
+	}
+	if db.AutoDumpSchema {
+		_ = db.DumpSchema()
+	}
+	return nil
+}
+
+// RollbackOnly rolls back exactly one applied migration that matches the given version.
+func (db *DB) RollbackOnly(version string) error {
+	drv, err := db.Driver()
+	if err != nil {
+		return err
+	}
+
+	migrations, err := db.FindMigrations()
+	if err != nil {
+		return err
+	}
+
+	var target *Migration
+	for i, m := range migrations {
+		if m.Version == version {
+			target = &migrations[i]
+			break
+		}
+	}
+	if target == nil {
+		return ErrMigrationNotFound
+	}
+	if !target.Applied {
+		return fmt.Errorf("migration `%s` is not applied, nothing to rollback", version)
+	}
+
+	sqlDB, err := db.openDatabaseForMigration(drv)
+	if err != nil {
+		return err
+	}
+	defer dbutil.MustClose(sqlDB)
+
+	fmt.Fprintf(db.Log, "Rolling back: %s\n", target.FileName)
+	start := time.Now()
+
+	parsed, err := target.Parse()
+	if err != nil {
+		return err
+	}
+	exec := func(tx dbutil.Transaction) error {
+		res, err := tx.Exec(parsed.Down)
+		if err != nil {
+			return drv.QueryError(parsed.Down, err)
+		} else if db.Verbose {
+			db.printVerbose(res)
+		}
+		return drv.DeleteMigration(tx, target.Version)
+	}
+
+	if parsed.DownOptions.Transaction() {
+		err = doTransaction(sqlDB, exec)
+	} else {
+		err = exec(sqlDB)
+	}
+
+	fmt.Fprintf(db.Log, "Rolled back: %s in %s\n", target.FileName, time.Since(start))
+	if err != nil {
+		return err
+	}
+	if db.AutoDumpSchema {
+		_ = db.DumpSchema()
+	}
+	return nil
+}
+
 // Status shows the status of all migrations
 func (db *DB) Status(quiet bool) (int, error) {
 	results, err := db.FindMigrations()

pkg/dbmate/db_test.go

@@ -593,6 +593,87 @@ func TestRollbackAll(t *testing.T) {
 	})
 }
 
+func TestMigrateOnly(t *testing.T) {
+	testEachURL(t, func(t *testing.T, u *url.URL) {
+		const v1 = "20151129054053" // users
+		const v2 = "20200227231541" // posts
+
+		db := newTestDB(t, u)
+		drv, err := db.Driver()
+		require.NoError(t, err)
+
+		require.NoError(t, db.Drop())
+		require.NoError(t, db.Create())
+
+		require.NoError(t, db.MigrateOnly(v1))
+
+		sqlDB, err := drv.Open()
+		require.NoError(t, err)
+		defer dbutil.MustClose(sqlDB)
+
+		applied, err := drv.SelectMigrations(sqlDB, -1)
+		require.NoError(t, err)
+		require.Equal(t, map[string]bool{v1: true}, applied)
+
+		var cnt int
+		require.NoError(t, sqlDB.QueryRow("select count(*) from users").Scan(&cnt))
+		require.Error(t, sqlDB.QueryRow("select count(*) from posts").Scan(&cnt))
+
+		require.NoError(t, db.MigrateOnly(v1))
+
+		err = db.MigrateOnly("99999999999999")
+		require.ErrorIs(t, err, dbmate.ErrMigrationNotFound)
+
+		db.Strict = true
+		err = db.MigrateOnly(v2)
+		require.NoError(t, err)
+
+		require.NoError(t, db.MigrateOnly(v1))
+	})
+}
+
+func TestRollbackOnly(t *testing.T) {
+	testEachURL(t, func(t *testing.T, u *url.URL) {
+		const v1 = "20151129054053" // users
+		const v2 = "20200227231541" // posts
+
+		db := newTestDB(t, u)
+		drv, err := db.Driver()
+		require.NoError(t, err)
+
+		require.NoError(t, db.Drop())
+		require.NoError(t, db.Create())
+		require.NoError(t, db.Migrate())
+
+		sqlDB, err := drv.Open()
+		require.NoError(t, err)
+		defer dbutil.MustClose(sqlDB)
+
+		applied, _ := drv.SelectMigrations(sqlDB, -1)
+		require.Equal(t, map[string]bool{v1: true, v2: true}, applied)
+
+		require.NoError(t, db.RollbackOnly(v2))
+
+		applied, _ = drv.SelectMigrations(sqlDB, -1)
+		require.Equal(t, map[string]bool{v1: true}, applied)
+
+		var cnt int
+		require.NoError(t, sqlDB.QueryRow("select count(*) from users").Scan(&cnt))
+		require.Error(t, sqlDB.QueryRow("select count(*) from posts").Scan(&cnt))
+
+		err = db.RollbackOnly(v2)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "not applied")
+
+		err = db.RollbackOnly("99999999999999")
+		require.ErrorIs(t, err, dbmate.ErrMigrationNotFound)
+
+		require.NoError(t, db.RollbackOnly(v1))
+		applied, _ = drv.SelectMigrations(sqlDB, -1)
+		require.Empty(t, applied)
+	})
+}
+
 func TestFindMigrations(t *testing.T) {
 	testEachURL(t, func(t *testing.T, u *url.URL) {
 		db := newTestDB(t, u)