Use MigrationKeeper

cosmos · Feb 10, 2021 · 41e0339 · 41e0339
1 parent 29b85b7
commit 41e0339
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 26 deletions.
diff --git a/simapp/app_test.go b/simapp/app_test.go
@@ -11,7 +11,6 @@ import (
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 	dbm "github.com/tendermint/tm-db"
 
-	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 )
@@ -68,11 +67,6 @@ func TestRunMigrations(t *testing.T) {
 		expRunErrMsg string
 		expCalled    int
 	}{
-		{
-			"cannot register migration for non-existant module",
-			"foo", 1,
-			true, "store key for module foo not found: not found", false, "", 0,
-		},
 		{
 			"cannot register migration for version 0",
 			"bank", 0,
@@ -107,7 +101,7 @@ func TestRunMigrations(t *testing.T) {
 
 			if tc.moduleName != "" {
 				// Register migration for module from version `forVersion` to `forVersion+1`.
-				err = app.configurator.RegisterMigration(tc.moduleName, tc.forVersion, func(sdk.Context, sdk.StoreKey, codec.Marshaler) error {
+				err = app.configurator.RegisterMigration(tc.moduleName, tc.forVersion, func(sdk.Context) error {
 					called++
 
 					return nil

diff --git a/types/module/configurator.go b/types/module/configurator.go
@@ -33,8 +33,6 @@ type configurator struct {
 	msgServer   grpc.Server
 	queryServer grpc.Server
 
-	// storeKeys is used to access module stores inside in-place store migrations.
-	storeKeys map[string]*sdk.KVStoreKey
 	// migrations is a map of moduleName -> forVersion -> migration script handler
 	migrations map[string]map[uint64]MigrationHandler
 }
@@ -45,7 +43,6 @@ func NewConfigurator(cdc codec.Marshaler, msgServer grpc.Server, queryServer grp
 		cdc:         cdc,
 		msgServer:   msgServer,
 		queryServer: queryServer,
-		storeKeys:   storeKeys,
 		migrations:  map[string]map[uint64]MigrationHandler{},
 	}
 }
@@ -76,11 +73,6 @@ func (c configurator) RegisterMigration(moduleName string, forVersion uint64, ha
 		return sdkerrors.Wrapf(sdkerrors.ErrLogic, "another migration for module %s and version %d already exists", moduleName, forVersion)
 	}
 
-	_, found := c.storeKeys[moduleName]
-	if !found {
-		return sdkerrors.Wrapf(sdkerrors.ErrNotFound, "store key for module %s not found", moduleName)
-	}
-
 	c.migrations[moduleName][forVersion] = handler
 
 	return nil
@@ -97,11 +89,6 @@ func (c configurator) runModuleMigrations(ctx sdk.Context, moduleName string, fr
 		return nil
 	}
 
-	storeKey, found := c.storeKeys[moduleName]
-	if !found {
-		return sdkerrors.Wrapf(sdkerrors.ErrNotFound, "store key for module %s not found", moduleName)
-	}
-
 	moduleMigrationsMap, found := c.migrations[moduleName]
 	if !found {
 		return sdkerrors.Wrapf(sdkerrors.ErrNotFound, "no migrations found for module %s", moduleName)
@@ -114,7 +101,7 @@ func (c configurator) runModuleMigrations(ctx sdk.Context, moduleName string, fr
 			return sdkerrors.Wrapf(sdkerrors.ErrNotFound, "no migration found for module %s from version %d to version %d", moduleName, i, i+1)
 		}
 
-		err := migrateFn(ctx, storeKey, c.cdc)
+		err := migrateFn(ctx)
 		if err != nil {
 			return err
 		}

diff --git a/types/module/module.go b/types/module/module.go
@@ -342,7 +342,7 @@ func (m *Manager) ExportGenesis(ctx sdk.Context, cdc codec.JSONMarshaler) map[st
 }
 
 // MigrationHandler is the migration function that each module registers.
-type MigrationHandler func(store sdk.Context, storeKey sdk.StoreKey, cdc codec.Marshaler) error
+type MigrationHandler func(store sdk.Context) error
 
 // MigrationMap is a map of moduleName -> version, where version denotes the
 // version from which we should perform the migration for each module.

diff --git a/x/bank/keeper/migrations.go b/x/bank/keeper/migrations.go
@@ -0,0 +1,20 @@
+package keeper
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	v042 "github.com/cosmos/cosmos-sdk/x/bank/legacy/v042"
+)
+
+// MigrationKeeper is an interface that the keeper implements for handling
+// in-place store migrations.
+type MigrationKeeper interface {
+	// Migrate1 migrates the store from version 1 to 2.
+	Migrate1(ctx sdk.Context) error
+}
+
+var _ MigrationKeeper = (*BaseKeeper)(nil)
+
+// Migrate1 implements MigrationKeeper.Migrate1 method.
+func (keeper BaseKeeper) Migrate1(ctx sdk.Context) error {
+	return v042.MigrateStore(ctx, keeper.storeKey)
+}
diff --git a/x/bank/legacy/v042/store.go b/x/bank/legacy/v042/store.go
@@ -1,7 +1,6 @@
 package v042
 
 import (
-	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	v040auth "github.com/cosmos/cosmos-sdk/x/auth/legacy/v040"
@@ -13,7 +12,7 @@ import (
 //
 // - Change addresses to be length-prefixed.
 // - Change balances prefix to 1 byte
-func MigrateStore(ctx sdk.Context, storeKey sdk.StoreKey, _ codec.Marshaler) error {
+func MigrateStore(ctx sdk.Context, storeKey sdk.StoreKey) error {
 	store := ctx.KVStore(storeKey)
 
 	// old key is of format:

diff --git a/x/bank/module.go b/x/bank/module.go
@@ -22,7 +22,6 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/bank/client/cli"
 	"github.com/cosmos/cosmos-sdk/x/bank/client/rest"
 	"github.com/cosmos/cosmos-sdk/x/bank/keeper"
-	v042 "github.com/cosmos/cosmos-sdk/x/bank/legacy/v042"
 	"github.com/cosmos/cosmos-sdk/x/bank/simulation"
 	"github.com/cosmos/cosmos-sdk/x/bank/types"
 )
@@ -101,7 +100,9 @@ type AppModule struct {
 func (am AppModule) RegisterServices(cfg module.Configurator) {
 	types.RegisterMsgServer(cfg.MsgServer(), keeper.NewMsgServerImpl(am.keeper))
 	types.RegisterQueryServer(cfg.QueryServer(), am.keeper)
-	cfg.RegisterMigration(types.ModuleName, 0, v042.MigrateStore)
+	cfg.RegisterMigration(types.ModuleName, 0, func(ctx sdk.Context) error {
+		return am.keeper.(keeper.MigrationKeeper).Migrate1(ctx)
+	})
 }
 
 // NewAppModule creates a new AppModule object