feat!: add migrator to the commit abci call

baseapp/abci.go

@@ -130,7 +130,7 @@ func (app *BaseApp) SetOption(req abci.RequestSetOption) (res abci.ResponseSetOp
 func (app *BaseApp) Info(req abci.RequestInfo) abci.ResponseInfo {
 	lastCommitID := app.cms.LastCommitID()
 	// load the app version for a non zero height and zero app hash
-	if lastCommitID.Version > 0 && app.appVersion == 0 {
+	if lastCommitID.Version > 0 {
 		ctx, err := app.createQueryContext(lastCommitID.Version, false)
 		if err != nil {
 			panic(err)
@@ -343,6 +343,36 @@ func (app *BaseApp) Commit() abci.ResponseCommit {
 	// The write to the DeliverTx state writes all state transitions to the root
 	// MultiStore (app.cms) so when Commit() is called is persists those values.
 	app.deliverState.ms.Write()
+
+	// Check if there has been an app version change. If so and there is a migrator
+	// set, begin to run migrations. This needs to be done before the commit so
+	// that the migrations are part of the app hash
+	if header.Version.App < app.appVersion &&
+		app.migrator.storeMigrator != nil &&
+		app.migrator.moduleMigrator != nil {
+
+		// first update the stores themselves by adding and removing them as necessary
+		storeMigrations, err := app.migrator.storeMigrator(header.Version.App, app.appVersion)
+		if err != nil {
+			panic(fmt.Sprintf("failed to get store migrations: %v", err))
+		}
+		app.MountKVStores(storeMigrations.Added)
+		err = app.cms.LoadLatestVersionAndUpgrade(storeMigrations.ToStoreUpgrades())
+		if err != nil {
+			panic(fmt.Sprintf("failed to upgrade stores: %v", err))
+		}
+
+		// create a new cached branch of the store to apply migrations to
+		app.setDeliverState(header)
+		err = app.migrator.moduleMigrator(app.deliverState.ctx, header.Version.App, app.appVersion)
+		if err != nil {
+			panic(fmt.Sprintf("failed to migrate modules: %v", err))
+		}
+
+		// write the new state to the branch
+		app.deliverState.ms.Write()
+	}
+
 	commitID := app.cms.Commit()
 
 	res := abci.ResponseCommit{

baseapp/baseapp.go

@@ -39,6 +39,21 @@ type (
 	// an older version of the software. In particular, if a module changed the substore key name
 	// (or removed a substore) between two versions of the software.
 	StoreLoader func(ms sdk.CommitMultiStore) error
+
+	// MigrateModuleFn gets called when there is a change in the app version. It is
+	// triggered in the Commit stage after all state transitions have occurred.
+	MigrateModuleFn func(ctx sdk.Context, fromVersion, toVersion uint64) error
+
+	// MigrateStoreFn gets called when there is a change in the app version. It is
+	// triggered in the Commit stage after all state transitions have occurred.
+	MigrateStoreFn func(fromVersion, toVersion uint64) (StoreMigrations, error)
+
+	// StoreMigrations is a type that contains the added and removed stores for a
+	// migration.
+	StoreMigrations struct {
+		Added   map[string]*storetypes.KVStoreKey
+		Deleted map[string]*storetypes.KVStoreKey
+	}
 )
 
 // BaseApp reflects the ABCI application implementation.
@@ -57,6 +72,7 @@ type BaseApp struct { // nolint: maligned
 	snapshotData
 	abciData
 	moduleRouter
+	migrator
 
 	// volatile states:
 	//
@@ -135,6 +151,11 @@ type appStore struct {
 	fauxMerkleMode bool // if true, IAVL MountStores uses MountStoresDB for simulation speed.
 }
 
+type migrator struct {
+	moduleMigrator MigrateModuleFn
+	storeMigrator  MigrateStoreFn
+}
+
 type moduleRouter struct {
 	router           sdk.Router        // handle any kind of message
 	queryRouter      sdk.QueryRouter   // router for redirecting query calls

baseapp/options.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/snapshots"
 	snapshottypes "github.com/cosmos/cosmos-sdk/snapshots/types"
 	"github.com/cosmos/cosmos-sdk/store"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -160,6 +161,20 @@ func (app *BaseApp) SetEndBlocker(endBlocker sdk.EndBlocker) {
 	app.endBlocker = endBlocker
 }
 
+func (app *BaseApp) SetMigrateModuleFn(migrator MigrateModuleFn) {
+	if app.sealed {
+		panic("cannot set migrate module fn: baseapp is sealed")
+	}
+	app.migrator.moduleMigrator = migrator
+}
+
+func (app *BaseApp) SetMigrateStoreFn(migrator MigrateStoreFn) {
+	if app.sealed {
+		panic("cannot set migrate store fn: baseapp is sealed")
+	}
+	app.migrator.storeMigrator = migrator
+}
+
 func (app *BaseApp) SetAnteHandler(ah sdk.AnteHandler) {
 	if app.sealed {
 		panic("SetAnteHandler() on sealed BaseApp")
@@ -263,3 +278,23 @@ func (app *BaseApp) SetQueryMultiStore(ms sdk.MultiStore) {
 	}
 	app.qms = ms
 }
+
+// ToStoreUpgrades converts the StoreMigrations to StoreUpgrades.
+func (sm StoreMigrations) ToStoreUpgrades() *storetypes.StoreUpgrades {
+	added := make([]string, len(sm.Added))
+	deleted := make([]string, len(sm.Deleted))
+	i := 0
+	for name := range sm.Added {
+		added[i] = name
+		i++
+	}
+	i = 0
+	for name := range sm.Deleted {
+		deleted[i] = name
+		i++
+	}
+	return &storetypes.StoreUpgrades{
+		Added:   added,
+		Deleted: deleted,
+	}
+}