Upgrades and Store Migrations

import (
    "github.com/cosmos/cosmos-sdk/x/upgrade"
    upgradekeeper "github.com/cosmos/cosmos-sdk/x/upgrade/keeper"
    upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
)

// Reserve a store key for x/upgrade alongside the other module keys:
// upgradetypes.StoreKey

// Create the UpgradeKeeper before the module manager. homePath is where
// ReadUpgradeInfoFromDisk looks for the upgrade-info.json a node writes at the
// halt height. The authority is usually the governance module account.
app.UpgradeKeeper = upgradekeeper.NewKeeper(
    skipUpgradeHeights,
    runtime.NewKVStoreService(keys[upgradetypes.StoreKey]),
    appCodec,
    homePath,
    app.BaseApp,
    authtypes.NewModuleAddress(govtypes.ModuleName).String(),
)

// Register the upgrade module with the module manager.
app.ModuleManager = module.NewManager(
    // other modules...
    upgrade.NewAppModule(app.UpgradeKeeper, app.AccountKeeper.AddressCodec()),
)

// Run x/upgrade first in PreBlock so it can detect a due plan and execute the
// handler before the rest of the block.
app.ModuleManager.SetOrderPreBlockers(
    upgradetypes.ModuleName,
    // other pre-blockers...
)
app.SetPreBlocker(app.PreBlocker)

func (app *MyApp) PreBlocker(ctx sdk.Context, _ *abci.RequestFinalizeBlock) (*sdk.ResponsePreBlock, error) {
    return app.ModuleManager.PreBlock(ctx)
}

func (app *MyApp) InitChainer(ctx sdk.Context, req *abci.RequestInitChain) (*abci.ResponseInitChain, error) {
    // ...
    if err := app.UpgradeKeeper.SetModuleVersionMap(ctx, app.ModuleManager.GetVersionMap()); err != nil {
        return nil, err
    }
    // ...
}

func (am AppModule) RegisterServices(cfg module.Configurator) {
    // --snip--
    if err := cfg.RegisterMigration(types.ModuleName, 1, func(ctx sdk.Context) error {
        // Perform in-place store migrations from ConsensusVersion 1 to 2.
        return nil
    }); err != nil {
        panic(fmt.Sprintf("failed to migrate %s from version 1 to 2: %v", types.ModuleName, err))
    }

    if err := cfg.RegisterMigration(types.ModuleName, 2, func(ctx sdk.Context) error {
        // Perform in-place store migrations from ConsensusVersion 2 to 3.
        return nil
    }); err != nil {
        panic(fmt.Sprintf("failed to migrate %s from version 2 to 3: %v", types.ModuleName, err))
    }
}

type Migrator struct {
    keeper *Keeper
}

func NewMigrator(keeper *Keeper) Migrator {
    return Migrator{keeper: keeper}
}

func (m Migrator) Migrate1to2(ctx sdk.Context) error {
    count, err := m.keeper.counter.Get(ctx)
    if err != nil {
        // A chain that never touched the counter has no stored value yet.
        if errors.Is(err, collections.ErrNotFound) {
            return nil
        }
        return err
    }

    return m.keeper.counter.Set(ctx, count*10)
}

package keeper

import (
    sdk "github.com/cosmos/cosmos-sdk/types"
    "github.com/cosmos/cosmos-sdk/x/bank/exported"
    v2 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v2"
    v3 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v3"
    v4 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v4"
)

// Migrator is a struct for handling in-place store migrations.
type Migrator struct {
    keeper         BaseKeeper
    legacySubspace exported.Subspace
}

// NewMigrator returns a new Migrator.
func NewMigrator(keeper BaseKeeper, legacySubspace exported.Subspace) Migrator {
    return Migrator{keeper: keeper, legacySubspace: legacySubspace}
}

// Migrate1to2 migrates from version 1 to 2.
func (m Migrator) Migrate1to2(ctx sdk.Context) error {
    return v2.MigrateStore(ctx, m.keeper.storeService, m.keeper.cdc)
}

// Migrate2to3 migrates x/bank storage from version 2 to 3.
func (m Migrator) Migrate2to3(ctx sdk.Context) error {
    return v3.MigrateStore(ctx, m.keeper.storeService, m.keeper.cdc)
}

// Migrate3to4 migrates x/bank storage from version 3 to 4.
func (m Migrator) Migrate3to4(ctx sdk.Context) error {
    m.MigrateSendEnabledParams(ctx)
    return v4.MigrateStore(ctx, m.keeper.storeService, m.legacySubspace, m.keeper.cdc)
}

type UpgradeHandler func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error)

app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
    // optional: additional upgrade logic here
    return app.ModuleManager.RunMigrations(ctx, app.Configurator(), fromVM)
})

upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
if err != nil {
    panic(err)
}
if upgradeInfo.Name == "my-plan" && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
    storeUpgrades := storetypes.StoreUpgrades{
        Added: []string{"newmodule"},
    }
    app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
}

fromVM["newmodule"] = newmodule.AppModule{}.ConsensusVersion()
return app.ModuleManager.RunMigrations(ctx, app.Configurator(), fromVM)

import foo "github.com/my/module/foo"

app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
    // Prevent RunMigrations from calling foo's default InitGenesis.
    fromVM["foo"] = foo.AppModule{}.ConsensusVersion()

    // Run your custom genesis initialization for foo.
    // InitGenesis takes sdk.Context, so unwrap the handler's context.Context.
    // myCustomGenesisState must be a json.RawMessage (the marshaled genesis state).
    app.ModuleManager.Modules["foo"].(module.HasGenesis).InitGenesis(sdk.UnwrapSDKContext(ctx), app.AppCodec(), myCustomGenesisState)

    return app.ModuleManager.RunMigrations(ctx, app.Configurator(), fromVM)
})