feat(runtime/v2): store loader on simappv2

runtime/v2/app.go

@@ -46,6 +46,8 @@ type App[T transaction.Tx] struct {
 	// GRPCMethodsToMessageMap maps gRPC method name to a function that decodes the request
 	// bytes into a gogoproto.Message, which then can be passed to appmanager.
 	GRPCMethodsToMessageMap map[string]func() gogoproto.Message
+
+	storeLoader StoreLoader
 }
 
 // Name returns the app name.
@@ -68,9 +70,14 @@ func (a *App[T]) DefaultGenesis() map[string]json.RawMessage {
 	return a.moduleManager.DefaultGenesis()
 }
 
+// SetStoreLoader sets the store loader.
+func (a *App[T]) SetStoreLoader(loader StoreLoader) {
+	a.storeLoader = loader
+}
+
 // LoadLatest loads the latest version.
 func (a *App[T]) LoadLatest() error {
-	return a.db.LoadLatestVersion()
+	return a.storeLoader(a.db)
 }
 
 // LoadHeight loads a particular height

runtime/v2/module.go

@@ -130,6 +130,7 @@ func ProvideAppBuilder[T transaction.Tx](
 		msgRouterBuilder:        msgRouterBuilder,
 		queryRouterBuilder:      stf.NewMsgRouterBuilder(), // TODO dedicated query router
 		GRPCMethodsToMessageMap: map[string]func() proto.Message{},
+		storeLoader:             DefaultStoreLoader,
 	}
 	appBuilder := &AppBuilder[T]{app: app}
 

runtime/v2/store.go

@@ -1,6 +1,8 @@
 package runtime
 
 import (
+	"fmt"
+
 	"cosmossdk.io/core/store"
 	"cosmossdk.io/server/v2/stf"
 	storev2 "cosmossdk.io/store/v2"
@@ -55,3 +57,32 @@ type Store interface {
 	// LastCommitID returns the latest commit ID
 	LastCommitID() (proof.CommitID, error)
 }
+
+// StoreLoader allows for custom loading of the store, this is useful when upgrading the store from a previous version
+type StoreLoader func(store Store) error
+
+// DefaultStoreLoader just calls LoadLatestVersion on the store
+func DefaultStoreLoader(store Store) error {
+	return store.LoadLatestVersion()
+}
+
+func V2UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades) StoreLoader {
+	return func(store Store) error {
+		latestVersion, err := store.GetLatestVersion()
+		if err != nil {
+			return err
+		}
+
+		if uint64(upgradeHeight) == latestVersion+1 {
+			if len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 {
+				if upgrader, ok := store.(storev2.UpgradeableStore); ok {
+					return upgrader.LoadVersionAndUpgrade(latestVersion, storeUpgrades)
+				}
+
+				return fmt.Errorf("store does not support upgrades")
+			}
+		}
+
+		return DefaultStoreLoader(store)
+	}
+}

simapp/v2/app_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/require"
 
+	app2 "cosmossdk.io/core/app"
 	"cosmossdk.io/core/comet"
 	context2 "cosmossdk.io/core/context"
 	"cosmossdk.io/core/server"
@@ -21,6 +22,7 @@ import (
 	serverv2 "cosmossdk.io/server/v2"
 	comettypes "cosmossdk.io/server/v2/cometbft/types"
 	"cosmossdk.io/store/v2/db"
+	authtypes "cosmossdk.io/x/auth/types"
 	banktypes "cosmossdk.io/x/bank/types"
 
 	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"

simapp/v2/upgrades.go

@@ -2,6 +2,7 @@ package simapp
 
 import (
 	"context"
+	"cosmossdk.io/runtime/v2"
 
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/core/store"
@@ -44,8 +45,6 @@ func (app *SimApp[T]) RegisterUpgradeHandlers() {
 			Deleted: []string{"crisis"}, // The SDK discontinued the crisis module in v0.52.0
 		}
 
-		// configure store loader that checks if version == upgradeHeight and applies store upgrades
-		_ = storeUpgrades
-		// app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
+		app.SetStoreLoader(runtime.V2UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
 	}
 }