feat(store/v2): implement the feature to upgrade the store keys

store/v2/commitment/store.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"sort"
 
 	protoio "github.com/cosmos/gogoproto/io"
 
@@ -120,7 +121,7 @@ func (c *CommitStore) GetLatestVersion() (uint64, error) {
 	return version, nil
 }
 
-func (c *CommitStore) LoadVersion(targetVersion uint64) error {
+func (c *CommitStore) LoadVersion(targetVersion uint64, upgrades *corestore.StoreUpgrades) error {
 	// Rollback the metadata to the target version.
 	latestVersion, err := c.GetLatestVersion()
 	if err != nil {
@@ -139,20 +140,43 @@ func (c *CommitStore) LoadVersion(targetVersion uint64) error {
 		}
 	}
 
-	for _, tree := range c.multiTrees {
+	storesKeys := make([]string, 0, len(c.multiTrees))
+	for storeKey := range c.multiTrees {
+		storesKeys = append(storesKeys, storeKey)
+	}
+
+	if upgrades != nil {
+		// deterministic iteration order for upgrades
+		// (as the underlying store may change and
+		// upgrades make store changes where the execution order may matter)
+		sort.Slice(storesKeys, func(i, j int) bool {
+			return storesKeys[i] < storesKeys[j]
+		})
+	}
+
+	for _, storeKey := range storesKeys {
+		tree := c.multiTrees[storeKey]
+
+		// If it has been added, set the initial version.
+		if upgrades.IsAdded(storeKey) || upgrades.RenamedFrom(storeKey) != "" {
+			if err := tree.SetInitialVersion(targetVersion + 1); err != nil {
+				return err
+			}
+		}
+
 		if err := tree.LoadVersion(targetVersion); err != nil {
 			return err
 		}
 	}
 
 	// If the target version is greater than the latest version, it is the snapshot
 	// restore case, we should create a new commit info for the target version.
-	var cInfo *proof.CommitInfo
 	if targetVersion > latestVersion {
-		cInfo = c.WorkingCommitInfo(targetVersion)
+		cInfo := c.WorkingCommitInfo(targetVersion)
+		return c.flushCommitInfo(targetVersion, cInfo)
 	}
 
-	return c.flushCommitInfo(targetVersion, cInfo)
+	return nil
 }
 
 func (c *CommitStore) GetCommitInfo(version uint64) (*proof.CommitInfo, error) {
@@ -470,7 +494,7 @@ loop:
 		}
 	}
 
-	return snapshotItem, c.LoadVersion(version)
+	return snapshotItem, c.LoadVersion(version, nil)
 }
 
 func (c *CommitStore) Close() (ferr error) {

store/v2/database.go

@@ -42,7 +42,7 @@ type Committer interface {
 	GetLatestVersion() (uint64, error)
 
 	// LoadVersion loads the tree at the given version.
-	LoadVersion(targetVersion uint64) error
+	LoadVersion(targetVersion uint64, upgrades *corestore.StoreUpgrades) error
 
 	// Commit commits the working tree to the database.
 	Commit(version uint64) (*proof.CommitInfo, error)
@@ -56,7 +56,7 @@ type Committer interface {
 	// Once migration is complete, this method should be removed and/or not used.
 	Get(storeKey []byte, version uint64, key []byte) ([]byte, error)
 
-	// SetInitialVersion sets the initial version of the tree.
+	// SetInitialVersion sets the initial version of the committer.
 	SetInitialVersion(version uint64) error
 
 	// GetCommitInfo returns the CommitInfo for the given version.

store/v2/root/store.go

@@ -215,7 +215,7 @@
 		return err
 	}
 
-	return s.loadVersion(lv)
+	return s.loadVersion(lv, nil)
 }
 
 func (s *Store) LoadVersion(version uint64) error {
@@ -224,13 +224,31 @@
 		defer s.telemetry.MeasureSince(now, "root_store", "load_version")
 	}
 
-	return s.loadVersion(version)
+	return s.loadVersion(version, nil)
 }
 
-func (s *Store) loadVersion(v uint64) error {
+// LoadVersionAndUpgrade implements the UpgradeableRootStore interface.
+//
+// NOTE: It cannot be called while the store is migrating.
+func (s *Store) LoadVersionAndUpgrade(version uint64, upgrades *corestore.StoreUpgrades) error {
+	if s.telemetry != nil {
+		now := time.Now()
+		defer s.telemetry.MeasureSince(now, "root_store", "load_version_and_upgrade")
+	}
+
+	if s.isMigrating {
+		return fmt.Errorf("cannot upgrade while migrating")
+	}
+
+	// TODO: Implement the pruning logic for the renamed/deleted store keys.
+
+	return s.loadVersion(version, upgrades)
+}
+
+func (s *Store) loadVersion(v uint64, upgrades *corestore.StoreUpgrades) error {
 	s.logger.Debug("loading version", "version", v)
 
-	if err := s.stateCommitment.LoadVersion(v); err != nil {
+	if err := s.stateCommitment.LoadVersion(v, upgrades); err != nil {
 		return fmt.Errorf("failed to load SC version %d: %w", v, err)
 	}