migrate: re-run post-step callbacks on error

This commit modifies the migration framework to re-attempt post-step
callbacks if they error during a migration. Previously, if a post-step
callbacks failed, but their associated SQL migration succeeded, the
post-step callback would not be re-attempted on the next migration run,
and instead proceed with the next SQL migration.

This is achieved by introducing the concept of a "post-step callback"
migration version. Post-step callbacks are their corresponding SQL
migration version offset by +1000000000.

During the execution of a post-step callback, the post-step callback
migration version will be persisted as the database version.
That way, if the post-step callback errors, the version for the database
will be the post-step callback version on the next startup. The
post-step callback will then be re-attempted before proceeding with
the next SQL migration.

Author: ViktorT-11

migrate.go

@@ -241,6 +241,23 @@ func (m *Migrate) Migrate(version uint) error {
 		return m.unlockErr(ErrDirty{curVersion})
 	}
 
+	// If the current version is a clean migration task version, then we
+	// need to rerun the task for the previous version before we can
+	// continue with any SQL migration(s).
+	if InTaskVersionRange(curVersion) {
+		sqlMigVersion := SQLMigrationVersion(curVersion)
+
+		err := m.execTaskAtMigVersion(sqlMigVersion)
+		if err != nil {
+			return m.unlockErr(err)
+		}
+
+		curVersion, dirty, err = m.databaseDrv.Version()
+		if err != nil {
+			return m.unlockErr(err)
+		}
+	}
+
 	ret := make(chan interface{}, m.PrefetchMigrations)
 	go m.read(curVersion, int(version), ret)
 
@@ -267,6 +284,23 @@ func (m *Migrate) Steps(n int) error {
 		return m.unlockErr(ErrDirty{curVersion})
 	}
 
+	// If the current version is a clean migration task version, then we
+	// need to rerun the task for the previous version before we can
+	// continue with any SQL migration(s).
+	if InTaskVersionRange(curVersion) {
+		sqlMigVersion := SQLMigrationVersion(curVersion)
+
+		err := m.execTaskAtMigVersion(sqlMigVersion)
+		if err != nil {
+			return m.unlockErr(err)
+		}
+
+		curVersion, dirty, err = m.databaseDrv.Version()
+		if err != nil {
+			return m.unlockErr(err)
+		}
+	}
+
 	ret := make(chan interface{}, m.PrefetchMigrations)
 
 	if n > 0 {
@@ -294,6 +328,23 @@ func (m *Migrate) Up() error {
 		return m.unlockErr(ErrDirty{curVersion})
 	}
 
+	// If the current version is a clean migration task version, then we
+	// need to rerun the task for the previous version before we can
+	// continue with any SQL migration(s).
+	if InTaskVersionRange(curVersion) {
+		sqlMigVersion := SQLMigrationVersion(curVersion)
+
+		err := m.execTaskAtMigVersion(sqlMigVersion)
+		if err != nil {
+			return m.unlockErr(err)
+		}
+
+		curVersion, dirty, err = m.databaseDrv.Version()
+		if err != nil {
+			return m.unlockErr(err)
+		}
+	}
+
 	ret := make(chan interface{}, m.PrefetchMigrations)
 
 	go m.readUp(curVersion, -1, ret)
@@ -316,6 +367,23 @@ func (m *Migrate) Down() error {
 		return m.unlockErr(ErrDirty{curVersion})
 	}
 
+	// If the current version is a clean migration task version, then we
+	// need to rerun the task for the previous version before we can
+	// continue with any SQL migration(s).
+	if InTaskVersionRange(curVersion) {
+		sqlMigVersion := SQLMigrationVersion(curVersion)
+
+		err := m.execTaskAtMigVersion(sqlMigVersion)
+		if err != nil {
+			return m.unlockErr(err)
+		}
+
+		curVersion, dirty, err = m.databaseDrv.Version()
+		if err != nil {
+			return m.unlockErr(err)
+		}
+	}
+
 	ret := make(chan interface{}, m.PrefetchMigrations)
 	go m.readDown(curVersion, -1, ret)
 	return m.unlockErr(m.runMigrations(ret))
@@ -354,6 +422,23 @@ func (m *Migrate) Run(migration ...*Migration) error {
 		return m.unlockErr(ErrDirty{curVersion})
 	}
 
+	// If the current version is a clean migration task version, then
+	// we need to rerun the task for the previous version before we can
+	// continue with any SQL migration(s).
+	if InTaskVersionRange(curVersion) {
+		sqlMigVersion := SQLMigrationVersion(curVersion)
+
+		err := m.execTaskAtMigVersion(sqlMigVersion)
+		if err != nil {
+			return m.unlockErr(err)
+		}
+
+		curVersion, dirty, err = m.databaseDrv.Version()
+		if err != nil {
+			return m.unlockErr(err)
+		}
+	}
+
 	ret := make(chan interface{}, m.PrefetchMigrations)
 
 	go func() {
@@ -732,6 +817,30 @@ func (m *Migrate) readDown(from int, limit int, ret chan<- interface{}) {
 	}
 }
 
+// readSingle reads a single migration for the given version, and sends it
+// over the passed channel.
+func (m *Migrate) readSingle(ver uint, ret chan<- interface{}) {
+	defer close(ret)
+
+	if err := m.versionExists(ver); err != nil {
+		ret <- err
+		return
+	}
+
+	migr, err := m.newMigration(ver, int(ver))
+	if err != nil {
+		ret <- err
+		return
+	}
+
+	ret <- migr
+	go func() {
+		if err := migr.Buffer(); err != nil {
+			m.logErr(err)
+		}
+	}()
+}
+
 // runMigrations reads *Migration and error from a channel. Any other type
 // sent on this channel will result in a panic. Each migration is then
 // proxied to the database driver and run against the database.
@@ -752,6 +861,12 @@ func (m *Migrate) runMigrations(ret <-chan interface{}) error {
 		case *Migration:
 			migr := r
 
+			if migr.Version >= TaskVersionOffset {
+				return fmt.Errorf("migration version %v is "+
+					"invalid, must be < %v", migr.Version,
+					TaskVersionOffset)
+			}
+
 			// set version with dirty state
 			if err := m.databaseDrv.SetVersion(migr.TargetVersion, true); err != nil {
 				return err
@@ -763,23 +878,9 @@ func (m *Migrate) runMigrations(ret <-chan interface{}) error {
 					return err
 				}
 
-				// If there is a task function for this
-				// migration, run it now.
-				cb, ok := m.opts.tasks[migr.Version]
-				if ok {
-					m.logVerbosePrintf("Running migration "+
-						"task for %v\n", migr.LogString())
-
-					err := cb(migr, m.databaseDrv)
-					if err != nil {
-						return fmt.Errorf("failed to "+
-							"execute migration "+
-							"task: %w",
-							err)
-					}
-
-					m.logVerbosePrintf("Migration task "+
-						"finished for %v\n", migr.LogString())
+				err := m.execTask(migr)
+				if err != nil {
+					return err
 				}
 			}
 
@@ -808,6 +909,108 @@ func (m *Migrate) runMigrations(ret <-chan interface{}) error {
 	return nil
 }
 
+// execTask checks if a migration task exists for the passed migration and
+// proceeds to execute if one exists.
+func (m *Migrate) execTask(migr *Migration) error {
+	task, ok := m.opts.tasks[migr.Version]
+	if ok {
+		m.logVerbosePrintf("Running migration task for %v\n",
+			migr.LogString())
+
+		taskVersion := int(migr.Version) + TaskVersionOffset
+
+		// Persist that we are in the migration task phase for this
+		// version.
+		if err := m.databaseDrv.SetVersion(taskVersion, true); err != nil {
+			return err
+		}
+
+		err := task(migr, m.databaseDrv)
+		if err != nil {
+			// Mark the database version as the taskVersion but in a
+			// clean state, to indicate that the migration task
+			// errored. We will therefore re-run the task on the
+			// next migration run.
+			if setErr := m.databaseDrv.SetVersion(taskVersion, false); setErr != nil {
+				// Note that if we error here, the database
+				// version will remain in a dirty state. As we
+				// cannot know if the migration task was
+				// executed or not in that scenario, manual
+				// intervention is required.
+				return fmt.Errorf("WARNING, failed to set "+
+					"migration version after migration "+
+					"task errored. Manual intervention "+
+					"needed! Migration task error: %w, "+
+					"version setting error : %w",
+					err, setErr)
+			}
+
+			return fmt.Errorf("failed to execute migration "+
+				"task: %w", err)
+		}
+
+		m.logVerbosePrintf("Migration task finished for %v\n",
+			migr.LogString())
+	}
+
+	return nil
+}
+
+// execTaskAtMigVersion executes only the migration task for the passed SQL
+// migration version.
+// The function can be used to re-execute the task for a SQL migration version
+// where the SQL migration was successfully applied, but where the task failed.
+func (m *Migrate) execTaskAtMigVersion(sqlMigVersion int) error {
+	var (
+		r      interface{}
+		migRet = make(chan interface{}, m.PrefetchMigrations)
+		err    error
+	)
+
+	// Fetch the migration for the specified SQL migration version.
+	go m.readSingle(uint(sqlMigVersion), migRet)
+
+	select {
+	case r = <-migRet:
+	case <-time.After(30 * time.Second):
+		return fmt.Errorf("timeout waiting for single migration "+
+			"version %v", sqlMigVersion)
+	}
+
+	if m.stop() {
+		return nil
+	}
+
+	switch r := r.(type) {
+	case *Migration:
+		// If the migration was found, execute the migration task.
+		migr := r
+
+		err = m.execTask(migr)
+		if err != nil {
+			return err
+		}
+
+		m.logVerbosePrintf("successfully re-executed migration task "+
+			"for SQL migration version: %v\n", sqlMigVersion)
+
+		// set clean state
+		if err = m.databaseDrv.SetVersion(migr.TargetVersion, false); err != nil {
+			return err
+		}
+
+		return nil
+
+	case error:
+		return fmt.Errorf("reading SQL migration at version "+
+			"%v failed: %w", sqlMigVersion, r)
+
+	default:
+		return fmt.Errorf("unknown type: %T when reading "+
+			"single migration", r)
+	}
+}
+
 // versionExists checks the source if either the up or down migration for
 // the specified migration version exists.
 func (m *Migrate) versionExists(version uint) (result error) {