From 42935909e771feeeee9a22d9268940337439d384 Mon Sep 17 00:00:00 2001
From: Uday Patil <udpatil@gmail.com>
Date: Fri, 9 Aug 2024 17:39:05 -0400
Subject: [PATCH] Delete kvstore specified in store upgrades

---
 store/rootmulti/store.go   | 13 +++++++++++++
 storev2/rootmulti/store.go | 13 +++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index 8a5fa496d..2ed486708 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -256,6 +256,15 @@ func (rs *Store) loadVersion(ver int64, upgrades *types.StoreUpgrades) error {
 		storesKeys = append(storesKeys, key)
 	}
 	if upgrades != nil {
+		for _, upgrade := range upgrades.Deleted {
+			deletionStoreKey := types.NewKVStoreKey(upgrade)
+			storesKeys = append(storesKeys, deletionStoreKey)
+			rs.storesParams[deletionStoreKey] = storeParams{
+				key: deletionStoreKey,
+				typ: types.StoreTypeIAVL, // TODO: is this safe
+			}
+			rs.keysByName[upgrade] = deletionStoreKey
+		}
 		// deterministic iteration order for upgrades
 		// (as the underlying store may change and
 		// upgrades make store changes where the execution order may matter)
@@ -283,6 +292,10 @@ func (rs *Store) loadVersion(ver int64, upgrades *types.StoreUpgrades) error {
 		// If it was deleted, remove all data
 		if upgrades.IsDeleted(key.Name()) {
 			deleteKVStore(store.(types.KVStore))
+			// drop deleted KV store from stores
+			delete(newStores, key)
+			delete(rs.keysByName, key.Name())
+			delete(rs.storesParams, key)
 		} else if oldName := upgrades.RenamedFrom(key.Name()); oldName != "" {
 			// handle renames specially
 			// make an unregistered key to satify loadCommitStore params
diff --git a/storev2/rootmulti/store.go b/storev2/rootmulti/store.go
index e4abfbf25..5b747db52 100644
--- a/storev2/rootmulti/store.go
+++ b/storev2/rootmulti/store.go
@@ -381,6 +381,14 @@ func (rs *Store) LoadVersionAndUpgrade(version int64, upgrades *types.StoreUpgra
 	for key := range rs.storesParams {
 		storesKeys = append(storesKeys, key)
 	}
+
+	if upgrades != nil {
+		// load storeKeys for deletion
+		for _, upgrade := range upgrades.Deleted {
+			deletionStoreKey := types.NewKVStoreKey(upgrade)
+			storesKeys = append(storesKeys, deletionStoreKey)
+		}
+	}
 	// deterministic iteration order for upgrades
 	sort.Slice(storesKeys, func(i, j int) bool {
 		return storesKeys[i].Name() < storesKeys[j].Name()
@@ -397,11 +405,13 @@ func (rs *Store) LoadVersionAndUpgrade(version int64, upgrades *types.StoreUpgra
 		return nil
 	}
 
+	storesKeysForDeletion := make(map[types.StoreKey]struct{})
 	var treeUpgrades []*proto.TreeNameUpgrade
 	for _, key := range storesKeys {
 		switch {
 		case upgrades.IsDeleted(key.Name()):
 			treeUpgrades = append(treeUpgrades, &proto.TreeNameUpgrade{Name: key.Name(), Delete: true})
+			storesKeysForDeletion[key] = struct{}{}
 		case upgrades.IsAdded(key.Name()) || upgrades.RenamedFrom(key.Name()) != "":
 			treeUpgrades = append(treeUpgrades, &proto.TreeNameUpgrade{Name: key.Name(), RenameFrom: upgrades.RenamedFrom(key.Name())})
 		}
@@ -415,6 +425,9 @@ func (rs *Store) LoadVersionAndUpgrade(version int64, upgrades *types.StoreUpgra
 	var err error
 	newStores := make(map[types.StoreKey]types.CommitKVStore, len(storesKeys))
 	for _, key := range storesKeys {
+		if _, ok := storesKeysForDeletion[key]; ok {
+			continue
+		}
 		newStores[key], err = rs.loadCommitStoreFromParams(key, rs.storesParams[key])
 		if err != nil {
 			return err