Fix upgrade store loader

store/rootmulti/store.go

@@ -315,7 +315,9 @@ func (rs *Store) TracingEnabled() bool {
 // LastCommitID implements Committer/CommitStore.
 func (rs *Store) LastCommitID() types.CommitID {
 	if rs.lastCommitInfo == nil {
-		return types.CommitID{}
+		return types.CommitID{
+			Version: getLatestVersion(rs.db),
+		}
 	}
 
 	return rs.lastCommitInfo.CommitID()

x/upgrade/types/storeloader.go

@@ -10,9 +10,9 @@ import (
 // pattern. This is useful for custom upgrade loading logic.
 func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades) baseapp.StoreLoader {
 	return func(ms sdk.CommitMultiStore) error {
-		if upgradeHeight == ms.LastCommitID().Version {
+		if upgradeHeight == ms.LastCommitID().Version+1 {
 			// Check if the current commit version and upgrade height matches
-			if len(storeUpgrades.Renamed) > 0 || len(storeUpgrades.Deleted) > 0 {
+			if len(storeUpgrades.Renamed) > 0 || len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 {
 				return ms.LoadLatestVersionAndUpgrade(storeUpgrades)
 			}
 		}

x/upgrade/types/storeloader_test.go

@@ -65,11 +65,13 @@ func checkStore(t *testing.T, db dbm.DB, ver int64, storeKey string, k, v []byte
 // Test that we can make commits and then reload old versions.
 // Test that LoadLatestVersion actually does.
 func TestSetLoader(t *testing.T) {
+	upgradeHeight := int64(5)
+
 	// set a temporary home dir
 	homeDir := t.TempDir()
 	upgradeInfoFilePath := filepath.Join(homeDir, "upgrade-info.json")
 	upgradeInfo := &store.UpgradeInfo{
-		Name: "test", Height: 0,
+		Name: "test", Height: upgradeHeight,
 	}
 
 	data, err := json.Marshal(upgradeInfo)
@@ -92,11 +94,12 @@ func TestSetLoader(t *testing.T) {
 			loadStoreKey: "foo",
 		},
 		"rename with inline opts": {
-			setLoader: useUpgradeLoader(0, &store.StoreUpgrades{
+			setLoader: useUpgradeLoader(upgradeHeight, &store.StoreUpgrades{
 				Renamed: []store.StoreRename{{
 					OldKey: "foo",
 					NewKey: "bar",
 				}},
+				Added: []string{"baz"},
 			}),
 			origStoreKey: "foo",
 			loadStoreKey: "bar",
@@ -116,25 +119,38 @@ func TestSetLoader(t *testing.T) {
 
 			// load the app with the existing db
 			opts := []func(*baseapp.BaseApp){baseapp.SetPruning(store.PruneNothing)}
+			capKey := sdk.NewKVStoreKey("main")
+
+			origapp := baseapp.NewBaseApp(t.Name(), defaultLogger(), db, nil, opts...)
+			origapp.MountStores(capKey)
+			origapp.MountStores(sdk.NewKVStoreKey(tc.loadStoreKey))
+			err := origapp.LoadLatestVersion()
+
+			for i := 2; i <= int(upgradeHeight-1); i++ {
+				origapp.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: int64(i)}})
+				res := origapp.Commit()
+				require.NotNil(t, res.Data)
+			}
+
 			if tc.setLoader != nil {
 				opts = append(opts, tc.setLoader)
 			}
 
+			// load the new app with the original app db
 			app := baseapp.NewBaseApp(t.Name(), defaultLogger(), db, nil, opts...)
-			capKey := sdk.NewKVStoreKey("main")
 			app.MountStores(capKey)
 			app.MountStores(sdk.NewKVStoreKey(tc.loadStoreKey))
-			err := app.LoadLatestVersion()
+			err = app.LoadLatestVersion()
 			require.Nil(t, err)
 
 			// "execute" one block
-			app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: 2}})
+			app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: upgradeHeight}})
 			res := app.Commit()
 			require.NotNil(t, res.Data)
 
 			// check db is properly updated
-			checkStore(t, db, 2, tc.loadStoreKey, k, v)
-			checkStore(t, db, 2, tc.loadStoreKey, []byte("foo"), nil)
+			checkStore(t, db, upgradeHeight, tc.loadStoreKey, k, v)
+			checkStore(t, db, upgradeHeight, tc.loadStoreKey, []byte("foo"), nil)
 		})
 	}
 }