From d1182f79fa0c5bbfab556537154457dcc3dae18c Mon Sep 17 00:00:00 2001
From: Alex Gartner <git@agartner.com>
Date: Thu, 14 Nov 2024 09:19:02 -0800
Subject: [PATCH] refactor: use major version in release upgrade handlers

---
 app/setup_handlers.go                         | 23 ++++++++++++++++---
 cmd/zetacored/root.go                         |  1 +
 cmd/zetacored/version.go                      | 18 +++++++++++++++
 .../scripts/start-upgrade-orchestrator.sh     |  2 +-
 4 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 cmd/zetacored/version.go

diff --git a/app/setup_handlers.go b/app/setup_handlers.go
index 04a94da668..7b609523e9 100644
--- a/app/setup_handlers.go
+++ b/app/setup_handlers.go
@@ -9,8 +9,23 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 
 	"github.com/zeta-chain/node/pkg/constant"
+	"golang.org/x/mod/semver"
 )
 
+// GetDefaultUpgradeHandlerVersion prints the default upgrade handler version
+//
+// There may be multiple upgrade handlers configured on some releases if different
+// migrations needto be run in different environment
+func GetDefaultUpgradeHandlerVersion() string {
+	// development builds always use the full version in the release handlers
+	if semver.Build(constant.Version) != "" || semver.Prerelease(constant.Version) != "" {
+		return constant.Version
+	}
+
+	// release builds use just the major version (v22.0.0 -> v22)
+	return semver.Major(constant.Version)
+}
+
 func SetupHandlers(app *App) {
 	allUpgrades := upgradeTracker{
 		upgrades: []upgradeTrackerItem{
@@ -50,10 +65,12 @@ func SetupHandlers(app *App) {
 		upgradeHandlerFns, storeUpgrades = allUpgrades.mergeAllUpgrades()
 	}
 
+	upgradeHandlerVersion := GetDefaultUpgradeHandlerVersion()
+
 	app.UpgradeKeeper.SetUpgradeHandler(
-		constant.Version,
+		upgradeHandlerVersion,
 		func(ctx sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
-			app.Logger().Info("Running upgrade handler for " + constant.Version)
+			app.Logger().Info("Running upgrade handler for " + upgradeHandlerVersion)
 
 			var err error
 			for _, upgradeHandler := range upgradeHandlerFns {
@@ -71,7 +88,7 @@ func SetupHandlers(app *App) {
 	if err != nil {
 		panic(err)
 	}
-	if upgradeInfo.Name == constant.Version && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
+	if upgradeInfo.Name == upgradeHandlerVersion && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
 		// Use upgrade store loader for the initial loading of all stores when app starts,
 		// it checks if version == upgradeHeight and applies store upgrades before loading the stores,
 		// so that new stores start with the correct version (the current height of chain),
diff --git a/cmd/zetacored/root.go b/cmd/zetacored/root.go
index 93ae4ba469..1f7115e57a 100644
--- a/cmd/zetacored/root.go
+++ b/cmd/zetacored/root.go
@@ -145,6 +145,7 @@ func initRootCmd(rootCmd *cobra.Command, encodingConfig types.EncodingConfig) {
 		GetPubKeyCmd(),
 		CollectObserverInfoCmd(),
 		AddrConversionCmd(),
+		UpgradeHandlerVersionCmd(),
 		tmcli.NewCompletionCmd(rootCmd, true),
 		ethermintclient.NewTestnetCmd(app.ModuleBasics, banktypes.GenesisBalancesIterator{}),
 
diff --git a/cmd/zetacored/version.go b/cmd/zetacored/version.go
new file mode 100644
index 0000000000..18555a5fa3
--- /dev/null
+++ b/cmd/zetacored/version.go
@@ -0,0 +1,18 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+	"github.com/zeta-chain/node/app"
+)
+
+func UpgradeHandlerVersionCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:   "upgrade-handler-version",
+		Short: "Print the upgrade handler version",
+		Run: func(cmd *cobra.Command, args []string) {
+			fmt.Println(app.GetDefaultUpgradeHandlerVersion())
+		},
+	}
+}
diff --git a/contrib/localnet/scripts/start-upgrade-orchestrator.sh b/contrib/localnet/scripts/start-upgrade-orchestrator.sh
index 1f8089bc3d..13d09a447f 100755
--- a/contrib/localnet/scripts/start-upgrade-orchestrator.sh
+++ b/contrib/localnet/scripts/start-upgrade-orchestrator.sh
@@ -39,7 +39,7 @@ fi
 # get new zetacored version
 curl -L -o /tmp/zetacored.new "${ZETACORED_URL}"
 chmod +x /tmp/zetacored.new
-UPGRADE_NAME=$(/tmp/zetacored.new version)
+UPGRADE_NAME=$(/tmp/zetacored.new upgrade-handler-version)
 
 # if explicit upgrade height not provided, use dumb estimator
 if [[ -z $UPGRADE_HEIGHT ]]; then