From 811d6a8c57cf7c583410b95a019d4098d4fb93df Mon Sep 17 00:00:00 2001
From: Shubham Sharma <samwodson@gmail.com>
Date: Tue, 1 Oct 2024 16:18:25 +0530
Subject: [PATCH] Add upgrade handler for "jip-2" and handle store upgrades

Introduced a new upgrade handler for "jip-2" with store upgrade functionalities. The handler now sets the store loader with the appropriate store upgrades and runs module migrations. Updated the app initialization to handle the new upgrade conditions.
---
 app/app.go     | 19 +++++++++++++++++--
 app/upgrade.go | 12 +++++++++---
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/app/app.go b/app/app.go
index 6aa7c8a..a68ecd9 100644
--- a/app/app.go
+++ b/app/app.go
@@ -1,6 +1,8 @@
 package app
 
 import (
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	trackgatemoduletypes "github.com/airchains-network/junction/x/trackgate/types"
 	"io"
 	"os"
 	"path/filepath"
@@ -354,15 +356,28 @@ func New(
 	// 	app.UpgradeKeeper.SetModuleVersionMap(ctx, app.ModuleManager.GetVersionMap())
 	// 	return app.App.InitChainer(ctx, req)
 	// })
+	// Handle store upgrades
+	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
+	if err != nil {
+		panic(err)
+	}
+
+	if upgradeInfo.Name == "jip-2" && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
+		storeUpgrades := storetypes.StoreUpgrades{
+			Added: []string{trackgatemoduletypes.StoreKey},
+		}
+
+		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
+	}
 
 	if err := app.Load(loadLatest); err != nil {
 		return nil, err
 	}
 
 	configurator := app.Configurator()
-	UpgradeHandleFunc := CreateDefaultUpgradeHandler(app.ModuleManager, configurator)
+	UpgradeHandleFunc := CreateDefaultUpgradeHandler(app.ModuleManager, configurator, app)
 	app.UpgradeKeeper.SetUpgradeHandler(
-		"jip-1",
+		"jip-2",
 		UpgradeHandleFunc, // Upgrade handler function
 	)
 
diff --git a/app/upgrade.go b/app/upgrade.go
index 74291b5..7082480 100644
--- a/app/upgrade.go
+++ b/app/upgrade.go
@@ -2,6 +2,8 @@ package app
 
 import (
 	"context"
+	storetypes "cosmossdk.io/store/types"
+	trackgatemoduletypes "github.com/airchains-network/junction/x/trackgate/types"
 
 	upgradetypes "cosmossdk.io/x/upgrade/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
@@ -11,12 +13,16 @@ import (
 func CreateDefaultUpgradeHandler(
 	mm *module.Manager,
 	configurator module.Configurator,
+	app *App,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
-		versionMap := module.VersionMap{
-			"junction": 100, // version 1.0.0
+		storeUpgrades := storetypes.StoreUpgrades{
+			Added: []string{trackgatemoduletypes.StoreKey},
 		}
 
-		return versionMap, nil
+		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(plan.Height, &storeUpgrades))
+
+		//return versionMap, nil
+		return app.ModuleManager.RunMigrations(ctx, configurator, fromVM)
 	}
 }