From 20c81092d36cd4225316f8c31ebb0bad3e8d0040 Mon Sep 17 00:00:00 2001
From: deepan95dev <deepan@vitwit.com>
Date: Mon, 26 Feb 2024 15:53:42 +0530
Subject: [PATCH] fix: upgrade handler with consensus, crisis

---
 app/app.go | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/app/app.go b/app/app.go
index 431c09f39..e633084d9 100644
--- a/app/app.go
+++ b/app/app.go
@@ -319,7 +319,6 @@ func New(
 		memKeys:           memKeys,
 		// invCheckPeriod:    invCheckPeriod,
 	}
-
 	app.ParamsKeeper = initParamsKeeper(appCodec, legacyAmino, keys[paramstypes.StoreKey], tkeys[paramstypes.TStoreKey])
 
 	// set the BaseApp's parameter store
@@ -447,7 +446,7 @@ func New(
 		app.BaseApp,
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
-
+	app.setupUpgradeStoreLoaders()
 	// Create IBC Keeper
 	app.IBCKeeper = ibckeeper.NewKeeper(
 		appCodec,
@@ -1160,6 +1159,26 @@ func BlockedAddresses() map[string]bool {
 	return modAccAddrs
 }
 
+// configure store loader that checks if version == upgradeHeight and applies store upgrades
+func (app *App) setupUpgradeStoreLoaders() {
+	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
+	if err != nil {
+		panic(fmt.Sprintf("failed to read upgrade info from disk %s", err))
+	}
+
+	if app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
+		return
+	}
+
+	storeUpgrades := storetypes.StoreUpgrades{
+		Added: []string{
+			consensusparamtypes.ModuleName,
+			crisistypes.ModuleName,
+		},
+	}
+	app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
+}
+
 func (app *App) RegisterUpgradeHandlers() {
 	baseAppLegacySS := app.ParamsKeeper.Subspace(baseapp.Paramspace).WithKeyTable(paramstypes.ConsensusParamsKeyTable())