feat: prepare app for upgrade (#106)

## Description

This PR adds an upgrade handler to make sure that the testnet can be
upgraded from the current `v1.0.0` version to the new `v1.1.0` version
properly.

It also includes some code fixes like missing autocli configurations and
services registering.

To make sure that upgrade goes smoothly, some migration implementations
have been added to properly set the params of new modules to default
values. Also, the params for the `x/assets` module has been removed as
it was empty and not used inside the codebase.

<!-- Add a description of the changes that this PR introduces and the
files that
are the most critical to review. -->

---

### Author Checklist

*All items are required. Please add a note to the item if the item is
not applicable and
please add links to any relevant follow up issues.*

I have...

- [ ] included the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] added `!` to the type prefix if API or client breaking change
- [ ] targeted the correct branch (see [PR
Targeting](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#pr-targeting))
- [ ] provided a link to the relevant issue or specification
- [ ] followed the guidelines for [building
modules](https://docs.cosmos.network/v0.44/building-modules/intro.html)
- [ ] included the necessary unit and integration
[tests](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#testing)
- [ ] added a changelog entry to `CHANGELOG.md`
- [ ] included comments for [documenting Go
code](https://blog.golang.org/godoc)
- [ ] updated the relevant documentation or specification
- [ ] reviewed "Files changed" and left comments if necessary
- [ ] confirmed all CI checks have passed

### Reviewers Checklist

*All items are required. Please add a note if the item is not applicable
and please add
your handle next to the items reviewed if you only reviewed selected
items.*

I have...

- [ ] confirmed the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] confirmed `!` in the type prefix if API or client breaking change
- [ ] confirmed all author checklist items have been addressed
- [ ] reviewed state machine logic
- [ ] reviewed API design and naming
- [ ] reviewed documentation is accurate
- [ ] reviewed tests and test coverage
- [ ] manually tested (if applicable)

---------

Co-authored-by: Hanjun Kim <hallazzang@gmail.com>

app/app.go

@@ -148,6 +148,7 @@ import (
 	apphook "github.com/milkyway-labs/milkyway/app/hook"
 	ibcwasmhooks "github.com/milkyway-labs/milkyway/app/ibc-hooks"
 	appkeepers "github.com/milkyway-labs/milkyway/app/keepers"
+	"github.com/milkyway-labs/milkyway/app/upgrades"
 	"github.com/milkyway-labs/milkyway/utils"
 	"github.com/milkyway-labs/milkyway/x/assets"
 	assetskeeper "github.com/milkyway-labs/milkyway/x/assets/keeper"
@@ -1093,7 +1094,7 @@ func NewMilkyWayApp(
 	app.indexerModule.RegisterServices(app.configurator)
 
 	// register upgrade handler for later use
-	app.RegisterUpgradeHandlers(app.configurator)
+	app.RegisterUpgradeHandlers()
 
 	autocliv1.RegisterQueryServer(app.GRPCQueryRouter(), runtimeservices.NewAutoCLIQueryService(app.ModuleManager.Modules))
 
@@ -1482,6 +1483,26 @@ func (app *MilkyWayApp) RegisterNodeService(clientCtx client.Context, cfg config
 	nodeservice.RegisterNodeService(clientCtx, app.GRPCQueryRouter(), cfg)
 }
 
+// Configurator returns the app's configurator.
+func (app *MilkyWayApp) Configurator() module.Configurator {
+	return app.configurator
+}
+
+// registerUpgrade registers the given upgrade to be supported by the app
+func (app *MilkyWayApp) registerUpgrade(upgrade upgrades.Upgrade) {
+	app.UpgradeKeeper.SetUpgradeHandler(upgrade.Name(), upgrade.Handler())
+
+	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
+	if err != nil {
+		panic(err)
+	}
+
+	if upgradeInfo.Name == upgrade.Name() && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
+		// Configure store loader that checks if version == upgradeHeight and applies store upgrades
+		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, upgrade.StoreUpgrades()))
+	}
+}
+
 // RegisterSwaggerAPI registers swagger route with API Server
 func RegisterSwaggerAPI(rtr *mux.Router) {
 	statikFS, err := fs.New()

app/upgrade.go

@@ -1,19 +1,10 @@
 package app
 
 import (
-	"context"
-
-	upgradetypes "cosmossdk.io/x/upgrade/types"
-	"github.com/cosmos/cosmos-sdk/types/module"
+	v110 "github.com/milkyway-labs/milkyway/app/upgrades/v110"
 )
 
-const upgradeName = "0.2.4"
-
 // RegisterUpgradeHandlers returns upgrade handlers
-func (app *MilkyWayApp) RegisterUpgradeHandlers(cfg module.Configurator) {
-	app.UpgradeKeeper.SetUpgradeHandler(upgradeName,
-		func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
-			return fromVM, nil
-		},
-	)
+func (app *MilkyWayApp) RegisterUpgradeHandlers() {
+	app.registerUpgrade(v110.NewUpgrade(app.ModuleManager, app.Configurator(), app.appCodec, app.keys, app.RewardsKeeper))
 }

app/upgrades/upgrade.go

@@ -0,0 +1,13 @@
+package upgrades
+
+import (
+	storetypes "cosmossdk.io/store/types"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+)
+
+// Upgrade represents a generic on-chain upgrade
+type Upgrade interface {
+	Name() string
+	Handler() upgradetypes.UpgradeHandler
+	StoreUpgrades() *storetypes.StoreUpgrades
+}

app/upgrades/v110/upgrade.go

@@ -0,0 +1,115 @@
+package v710
+
+import (
+	"context"
+
+	storetypes "cosmossdk.io/store/types"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	"github.com/cosmos/cosmos-sdk/codec"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	"github.com/milkyway-labs/milkyway/app/upgrades"
+	assetstypes "github.com/milkyway-labs/milkyway/x/assets/types"
+	operatorstypes "github.com/milkyway-labs/milkyway/x/operators/types"
+	poolstypes "github.com/milkyway-labs/milkyway/x/pools/types"
+	restakingtypes "github.com/milkyway-labs/milkyway/x/restaking/types"
+	rewardskeeper "github.com/milkyway-labs/milkyway/x/rewards/keeper"
+	rewardstypes "github.com/milkyway-labs/milkyway/x/rewards/types"
+	servicestypes "github.com/milkyway-labs/milkyway/x/services/types"
+)
+
+var (
+	_ upgrades.Upgrade = &Upgrade{}
+)
+
+// Upgrade represents the v1.1.0 upgrade
+type Upgrade struct {
+	mm           *module.Manager
+	configurator module.Configurator
+
+	cdc           codec.BinaryCodec
+	keys          map[string]*storetypes.KVStoreKey
+	rewardsKeeper *rewardskeeper.Keeper
+}
+
+// NewUpgrade returns a new Upgrade instance
+func NewUpgrade(
+	mm *module.Manager,
+	configurator module.Configurator,
+	cdc codec.BinaryCodec,
+	keys map[string]*storetypes.KVStoreKey,
+	rewardsKeeper *rewardskeeper.Keeper,
+) *Upgrade {
+	return &Upgrade{
+		mm:            mm,
+		configurator:  configurator,
+		cdc:           cdc,
+		keys:          keys,
+		rewardsKeeper: rewardsKeeper,
+	}
+}
+
+// Name implements upgrades.Upgrade
+func (u *Upgrade) Name() string {
+	return "v1.1.0"
+}
+
+// Handler implements upgrades.Upgrade
+func (u *Upgrade) Handler() upgradetypes.UpgradeHandler {
+	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+		// Set the operators module params
+		store := sdkCtx.KVStore(u.keys[operatorstypes.ModuleName])
+		operatorsParams := operatorstypes.DefaultParams()
+		store.Set(operatorstypes.ParamsKey, u.cdc.MustMarshal(&operatorsParams))
+
+		// Set the pools params
+		store = sdkCtx.KVStore(u.keys[poolstypes.ModuleName])
+		poolsParams := poolstypes.DefaultParams()
+		store.Set(poolstypes.ParamsKey, u.cdc.MustMarshal(&poolsParams))
+
+		// Set the restaking params
+		store = sdkCtx.KVStore(u.keys[restakingtypes.ModuleName])
+		restakingParams := restakingtypes.DefaultParams()
+		store.Set(restakingtypes.ParamsKey, u.cdc.MustMarshal(&restakingParams))
+
+		// Set the rewards params
+		if err := u.rewardsKeeper.Params.Set(ctx, rewardstypes.DefaultParams()); err != nil {
+			return nil, err
+		}
+
+		// Set the services params
+		store = sdkCtx.KVStore(u.keys[servicestypes.ModuleName])
+		servicesParams := servicestypes.DefaultParams()
+		store.Set(servicestypes.ParamsKey, u.cdc.MustMarshal(&servicesParams))
+
+		// Set the module versions
+		fromVM[assetstypes.ModuleName] = 1
+		fromVM[operatorstypes.ModuleName] = 1
+		fromVM[poolstypes.ModuleName] = 1
+		fromVM[restakingtypes.ModuleName] = 1
+		fromVM[rewardstypes.ModuleName] = 1
+		fromVM[servicestypes.ModuleName] = 1
+
+		// This upgrade does not require any migration, so we can simply return the current version map
+		return u.mm.RunMigrations(ctx, u.configurator, fromVM)
+	}
+}
+
+// StoreUpgrades implements upgrades.Upgrade
+func (u *Upgrade) StoreUpgrades() *storetypes.StoreUpgrades {
+	return &storetypes.StoreUpgrades{
+		Added: []string{
+			assetstypes.ModuleName,
+			operatorstypes.ModuleName,
+			poolstypes.ModuleName,
+			restakingtypes.ModuleName,
+			rewardstypes.ModuleName,
+			servicestypes.ModuleName,
+		},
+		Renamed: nil,
+		Deleted: nil,
+	}
+}

proto/milkyway/assets/v1/genesis.proto

@@ -3,17 +3,13 @@ package milkyway.assets.v1;
 
 import "amino/amino.proto";
 import "gogoproto/gogo.proto";
-import "milkyway/assets/v1/params.proto";
 import "milkyway/assets/v1/models.proto";
 
 option go_package = "github.com/milkyway-labs/milkyway/x/assets/types";
 
 // GenesisState defines the module's genesis state.
 message GenesisState {
-  // Params defines the parameters of the module.
-  Params params = 1 [ (gogoproto.nullable) = false ];
-
   // Assets defines the registered assets.
-  repeated Asset assets = 2
+  repeated Asset assets = 1
       [ (gogoproto.nullable) = false, (amino.dont_omitempty) = true ];
 }

proto/milkyway/assets/v1/messages.proto

@@ -5,7 +5,6 @@ import "amino/amino.proto";
 import "cosmos_proto/cosmos.proto";
 import "cosmos/msg/v1/msg.proto";
 import "gogoproto/gogo.proto";
-import "milkyway/assets/v1/params.proto";
 import "milkyway/assets/v1/models.proto";
 
 option go_package = "github.com/milkyway-labs/milkyway/x/assets/types";
@@ -20,11 +19,6 @@ service Msg {
   // DeregisterAsset defines the operation for de-registering an asset with
   // its denomination.
   rpc DeregisterAsset(MsgDeregisterAsset) returns (MsgDeregisterAssetResponse);
-
-  // UpdateParams defines a (governance) operation for updating the module
-  // parameters.
-  // The authority defaults to the x/gov module account.
-  rpc UpdateParams(MsgUpdateParams) returns (MsgUpdateParamsResponse);
 }
 
 // MsgRegisterAsset defines the message structure for the RegisterAsset
@@ -60,26 +54,4 @@ message MsgDeregisterAsset {
 }
 
 // MsgRegisterAssetResponse is the return value of MsgDeregisterAsset.
-message MsgDeregisterAssetResponse {}
-
-// MsgUpdateParams defines the message structure for the UpdateParams gRPC
-// service method. It allows the authority to update the module parameters.
-message MsgUpdateParams {
-  option (cosmos.msg.v1.signer) = "authority";
-  option (amino.name) = "milkyway/assets/MsgUpdateParams";
-
-  // Authority is the address that controls the module (defaults to x/gov unless
-  // overwritten).
-  string authority = 1 [
-    (gogoproto.moretags) = "yaml:\"authority\"",
-    (cosmos_proto.scalar) = "cosmos.AddressString"
-  ];
-
-  // Params define the parameters to update.
-  //
-  // NOTE: All parameters must be supplied.
-  Params params = 2 [ (gogoproto.nullable) = false ];
-}
-
-// MsgUpdateParamsResponse is the return value of MsgUpdateParams.
-message MsgUpdateParamsResponse {}
+message MsgDeregisterAssetResponse {}

proto/milkyway/assets/v1/query.proto

@@ -5,19 +5,12 @@ import "amino/amino.proto";
 import "gogoproto/gogo.proto";
 import "google/api/annotations.proto";
 import "cosmos/base/query/v1beta1/pagination.proto";
-import "milkyway/assets/v1/params.proto";
 import "milkyway/assets/v1/models.proto";
 
 option go_package = "github.com/milkyway-labs/milkyway/x/assets/types";
 
 // Query defines the gRPC querier service.
 service Query {
-  // Params defines a gRPC query method that returns the parameters of the
-  // module.
-  rpc Params(QueryParamsRequest) returns (QueryParamsResponse) {
-    option (google.api.http).get = "/milkyway/assets/v1/params";
-  }
-
   // Assets defined a gRPC query method that returns all assets registered.
   rpc Assets(QueryAssetsRequest) returns (QueryAssetsResponse) {
     option (google.api.http).get = "/milkyway/assets/v1/assets";
@@ -30,14 +23,6 @@ service Query {
   }
 }
 
-// QueryParamsRequest is the request type for the Query/Params RPC method.
-message QueryParamsRequest {}
-
-// QueryParamsResponse is the response type for the Query/Params RPC method.
-message QueryParamsResponse {
-  Params params = 1 [ (gogoproto.nullable) = false ];
-}
-
 // QueryAssetsRequest is the request type for the Query/Assets RPC method.
 message QueryAssetsRequest {
   // Ticker defines an optional filter parameter to query assets with the given

testutils/storetesting/context.go

@@ -0,0 +1,36 @@
+package storetesting
+
+import (
+	"cosmossdk.io/log"
+	"cosmossdk.io/store"
+	"cosmossdk.io/store/metrics"
+	storetypes "cosmossdk.io/store/types"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	dbm "github.com/cosmos/cosmos-db"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func BuildContext(
+	keys map[string]*storetypes.KVStoreKey, tKeys map[string]*storetypes.TransientStoreKey, memKeys map[string]*storetypes.MemoryStoreKey,
+) sdk.Context {
+	logger := log.NewNopLogger()
+	db := dbm.NewMemDB()
+	cms := store.NewCommitMultiStore(db, logger, metrics.NewNoOpMetrics())
+	for _, key := range keys {
+		cms.MountStoreWithDB(key, storetypes.StoreTypeIAVL, db)
+	}
+	for _, tKey := range tKeys {
+		cms.MountStoreWithDB(tKey, storetypes.StoreTypeTransient, db)
+	}
+	for _, memKey := range memKeys {
+		cms.MountStoreWithDB(memKey, storetypes.StoreTypeMemory, nil)
+	}
+
+	err := cms.LoadLatestVersion()
+	if err != nil {
+		panic(err)
+	}
+
+	return sdk.NewContext(cms, tmproto.Header{}, false, logger)
+}

x/assets/keeper/genesis.go

@@ -8,35 +8,23 @@ import (
 
 // ExportGenesis returns the GenesisState associated with the given context
 func (k *Keeper) ExportGenesis(ctx sdk.Context) (*types.GenesisState, error) {
-	// Get the params
-	params, err := k.Params.Get(ctx)
-	if err != nil {
-		return nil, err
-	}
-
 	// Get all the assets
 	var assets []types.Asset
 	_ = k.Assets.Walk(ctx, nil, func(_ string, asset types.Asset) (stop bool, err error) {
 		assets = append(assets, asset)
 		return false, nil
 	})
 
-	return types.NewGenesisState(params, assets), nil
+	return types.NewGenesisState(assets), nil
 }
 
 // --------------------------------------------------------------------------------------------------------------------
 
 // InitGenesis initializes the state from a GenesisState
 func (k *Keeper) InitGenesis(ctx sdk.Context, state *types.GenesisState) error {
-	// Store the params
-	err := k.Params.Set(ctx, state.Params)
-	if err != nil {
-		return err
-	}
-
 	// Store the assets
 	for _, asset := range state.Assets {
-		err = k.SetAsset(ctx, asset)
+		err := k.SetAsset(ctx, asset)
 		if err != nil {
 			return err
 		}