From 259bd524676b8e37775f3cf0b8a81c1295a34d4f Mon Sep 17 00:00:00 2001
From: Robert Zaremba <robert@zaremba.ch>
Date: Fri, 5 May 2023 19:42:56 +0200
Subject: [PATCH] feat: better upgrade registration

---
 app/upgrades.go | 61 +++++++++++++------------------------------------
 1 file changed, 16 insertions(+), 45 deletions(-)

diff --git a/app/upgrades.go b/app/upgrades.go
index d0b6993b20..38ff102b50 100644
--- a/app/upgrades.go
+++ b/app/upgrades.go
@@ -40,33 +40,16 @@ func (app UmeeApp) RegisterUpgradeHandlers(bool) {
 	}
 
 	app.registerUpgrade3_0(upgradeInfo)
-	app.registerUpgrade3_1(upgradeInfo)
+	app.registerUpgrade("v3.1.0", upgradeInfo)
 	app.registerUpgrade3_1to3_3(upgradeInfo)
 	app.registerUpgrade3_2to3_3(upgradeInfo)
 	app.registerUpgrade3_3to4_0(upgradeInfo)
-	app.registerUpgrade4_0_1(upgradeInfo)
+	app.registerUpgrade("v4.0.1", upgradeInfo)
 	app.registerUpgrade4_1(upgradeInfo)
-	app.registerUpgrade4_2(upgradeInfo)
+	app.registerUpgrade("v4.2", upgradeInfo, uibc.ModuleName)
 	app.registerUpgrade4_3(upgradeInfo)
-	app.registerUpgrade4_4(upgradeInfo)
-	app.registerUpgrade4_5(upgradeInfo)
-}
-
-// performs upgrade from v4.4 to v4.5
-func (app *UmeeApp) registerUpgrade4_5(upgradeInfo upgradetypes.Plan) {
-	const planName = "v4.5-alpha1" // TODO: set correct name
-	app.UpgradeKeeper.SetUpgradeHandler(planName, onlyModuleMigrations(app, planName))
-	app.storeUpgrade(planName, upgradeInfo, storetypes.StoreUpgrades{
-		Added: []string{
-			incentive.ModuleName,
-		},
-	})
-}
-
-// performs upgrade from v4.3 to v4.4
-func (app *UmeeApp) registerUpgrade4_4(_ upgradetypes.Plan) {
-	const planName = "v4.4"
-	app.UpgradeKeeper.SetUpgradeHandler(planName, onlyModuleMigrations(app, planName))
+	app.registerUpgrade("v4.4", upgradeInfo)
+	app.registerUpgrade("v4.5-alpha1", upgradeInfo, incentive.ModuleName) // TODO: set correct name
 }
 
 // performs upgrade from v4.2 to v4.3
@@ -119,17 +102,6 @@ func (app *UmeeApp) registerUpgrade4_3(upgradeInfo upgradetypes.Plan) {
 	})
 }
 
-// performs upgrade from v4.1 to v4.2
-func (app *UmeeApp) registerUpgrade4_2(upgradeInfo upgradetypes.Plan) {
-	const planName = "v4.2"
-	app.UpgradeKeeper.SetUpgradeHandler(planName, onlyModuleMigrations(app, planName))
-	app.storeUpgrade(planName, upgradeInfo, storetypes.StoreUpgrades{
-		Added: []string{
-			uibc.ModuleName,
-		},
-	})
-}
-
 // performs upgrade from v4.0 to v4.1
 func (app *UmeeApp) registerUpgrade4_1(_ upgradetypes.Plan) {
 	const planName = "v4.1.0"
@@ -152,12 +124,6 @@ func (app *UmeeApp) registerUpgrade4_1(_ upgradetypes.Plan) {
 	)
 }
 
-// performs upgrade from v4.0.0 to v4.0.1
-func (app *UmeeApp) registerUpgrade4_0_1(_ upgradetypes.Plan) {
-	const planName = "v4.0.1"
-	app.UpgradeKeeper.SetUpgradeHandler(planName, onlyModuleMigrations(app, planName))
-}
-
 // performs upgrade from v3.3 -> v4
 func (app *UmeeApp) registerUpgrade3_3to4_0(_ upgradetypes.Plan) {
 	const planName = "v4.0"
@@ -215,12 +181,6 @@ func (app *UmeeApp) registerUpgrade3_2to3_3(_ upgradetypes.Plan) {
 		})
 }
 
-// performs upgrade from v3.0 -> v3.1
-func (app *UmeeApp) registerUpgrade3_1(_ upgradetypes.Plan) {
-	const planName = "v3.1.0"
-	app.UpgradeKeeper.SetUpgradeHandler(planName, onlyModuleMigrations(app, planName))
-}
-
 // performs upgrade from v1->v3
 func (app *UmeeApp) registerUpgrade3_0(upgradeInfo upgradetypes.Plan) {
 	const planName = "v1.1-v3.0"
@@ -287,3 +247,14 @@ func (app *UmeeApp) storeUpgrade(planName string, ui upgradetypes.Plan, stores s
 			upgradetypes.UpgradeStoreLoader(ui.Height, &stores))
 	}
 }
+
+// registerUpgrade sets an upgrade handler which only runs module migrations
+// and adds new storages storages
+func (app *UmeeApp) registerUpgrade(planName string, upgradeInfo upgradetypes.Plan, newStores ...string) {
+	app.UpgradeKeeper.SetUpgradeHandler(planName, onlyModuleMigrations(app, planName))
+
+	if len(newStores) > 0 {
+		app.storeUpgrade(planName, upgradeInfo, storetypes.StoreUpgrades{
+			Added: newStores})
+	}
+}