feat: add --unsafe-remove-modules flag to Rollback cmd

server/rollback.go

@@ -8,11 +8,13 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/server/types"
+	"github.com/cosmos/cosmos-sdk/store/rootmulti"
 )
 
 // NewRollbackCmd creates a command to rollback tendermint and multistore state by one height.
 func NewRollbackCmd(appCreator types.AppCreator, defaultNodeHome string) *cobra.Command {
 	var removeBlock bool
+	moduleKeysToDelete := make([]string, 0)
 
 	cmd := &cobra.Command{
 		Use:   "rollback",
@@ -39,6 +41,17 @@ application.
 			if err != nil {
 				return fmt.Errorf("failed to rollback tendermint state: %w", err)
 			}
+
+			// for rolling back upgrades that add modules, we must first forcibly delete those.
+			// otherwise, the rollback will panic because no version of new modules will exist.
+			store := app.CommitMultiStore().(*rootmulti.Store)
+			for _, key := range moduleKeysToDelete {
+				fmt.Printf("deleting latest version of KVStore with key %s\n", key)
+				if err := store.DeleteLatestVersion(key); err != nil {
+					return err
+				}
+			}
+
 			// rollback the multistore
 
 			if err := app.CommitMultiStore().RollbackToVersion(height); err != nil {
@@ -52,5 +65,6 @@ application.
 
 	cmd.Flags().String(flags.FlagHome, defaultNodeHome, "The application home directory")
 	cmd.Flags().BoolVar(&removeBlock, "hard", false, "remove last block as well as state")
+	cmd.Flags().StringSliceVar(&moduleKeysToDelete, "unsafe-remove-modules", []string{}, "force delete KV stores with the provided prefix. useful for rolling back an upgrade that adds a module")
 	return cmd
 }

store/iavl/store.go

@@ -229,6 +229,11 @@ func (st *Store) DeleteVersionsTo(version int64) error {
 	return st.tree.DeleteVersionsTo(version)
 }
 
+// DeleteVersionsFrom deletes from the given version upwards form the MutableTree.
+func (st *Store) DeleteVersionsFrom(version int64) error {
+	return st.tree.DeleteVersionsFrom(version)
+}
+
 // LoadVersionForOverwriting attempts to load a tree at a previously committed
 // version, or the latest version below it. Any versions greater than targetVersion will be deleted.
 func (st *Store) LoadVersionForOverwriting(targetVersion int64) (int64, error) {

store/iavl/tree.go

@@ -30,6 +30,7 @@ type (
 		WorkingHash() []byte
 		VersionExists(version int64) bool
 		DeleteVersionsTo(version int64) error
+		DeleteVersionsFrom(version int64) error
 		GetVersioned(key []byte, version int64) ([]byte, error)
 		GetImmutable(version int64) (*iavl.ImmutableTree, error)
 		SetInitialVersion(version uint64)
@@ -67,6 +68,10 @@ func (it *immutableTree) DeleteVersionsTo(_ int64) error {
 	panic("cannot call 'DeleteVersionsTo' on an immutable IAVL tree")
 }
 
+func (it *immutableTree) DeleteVersionsFrom(_ int64) error {
+	panic("cannot call 'DeleteVersionsFrom' on an immutable IAVL tree")
+}
+
 func (it *immutableTree) SetInitialVersion(_ uint64) {
 	panic("cannot call 'SetInitialVersion' on an immutable IAVL tree")
 }

store/rootmulti/store.go

@@ -992,6 +992,68 @@ func (rs *Store) buildCommitInfo(version int64) *types.CommitInfo {
 	}
 }
 
+// DeleteLatestVersion finds a store with the given key name and deletes its latest version.
+// The store is deregistered from the rootmulti store. Calls to buildCommitInfo will not include it.
+// For stores with IAVL types, the deletion is written to the disk.
+// This is a destructive operation to be used with caution.
+// The reason it was added was to allow for rollbacks of upgrades that add modules.
+// Stores that do not exist in the version prior to upgrade can be forcibly deleted
+// before calling Rollback()
+func (rs *Store) DeleteLatestVersion(keyName string) error {
+	ver := GetLatestVersion(rs.db)
+	if ver == 0 {
+		return fmt.Errorf("unable to delete KVStore with key name %s, latest version is 0", keyName)
+	}
+
+	// Find the store key with the provided name
+	var key types.StoreKey = nil
+	for k := range rs.storesParams {
+		if k.Name() == keyName {
+			key = k
+			break
+		}
+	}
+	if key == nil {
+		return fmt.Errorf("no store found with key name %s", keyName)
+	}
+
+	// Get the KVStore for that key
+	cInfo, err := rs.GetCommitInfo(ver)
+	if err != nil {
+		return err
+	}
+	infos := make(map[string]types.StoreInfo)
+	for _, storeInfo := range cInfo.StoreInfos {
+		infos[storeInfo.Name] = storeInfo
+	}
+	commitID := rs.getCommitID(infos, key.Name())
+	store, err := rs.loadCommitStoreFromParams(key, commitID, rs.storesParams[key])
+	if err != nil {
+		return errors.Wrap(err, "failed to load store")
+	}
+
+	rs.logger.Debug("deleting KVStore", "key", key.Name(), "latest version", ver)
+
+	// for IAVL stores, commit the deletion of the latest version to disk.
+	if store.GetStoreType() == types.StoreTypeIAVL {
+		// unwrap the caching layer
+		store = rs.GetCommitKVStore(key)
+		if err := store.(*iavl.Store).DeleteVersionsFrom(ver); err != nil {
+			return errors.Wrapf(err, "failed to delete versions %d onwards of %s store", ver, key.Name())
+		}
+	}
+
+	// deregister store from the rootmulti store
+	// Any future buildCommitInfo will no longer include the store.
+	if _, ok := rs.stores[key]; ok {
+		delete(rs.stores, key)
+		delete(rs.storesParams, key)
+		delete(rs.keysByName, key.Name())
+	}
+
+	return nil
+}
+
 // RollbackToVersion delete the versions after `target` and update the latest version.
 func (rs *Store) RollbackToVersion(target int64) error {
 	if target <= 0 {
@@ -1003,9 +1065,10 @@ func (rs *Store) RollbackToVersion(target int64) error {
 			// If the store is wrapped with an inter-block cache, we must first unwrap
 			// it to get the underlying IAVL store.
 			store = rs.GetCommitKVStore(key)
+			rs.logger.Debug("loading version %d for store with key %s (%s)\n", target, key.String(), key.Name())
 			_, err := store.(*iavl.Store).LoadVersionForOverwriting(target)
 			if err != nil {
-				return err
+				return errors.Wrapf(err, "failed loading version %d for store with key name '%s'", target, key.Name())
 			}
 		}
 	}