storage snapshot api

cosmos · Nov 13, 2023 · c0dcf8f · c0dcf8f
1 parent 2569a3f
commit c0dcf8f
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 4 deletions.
diff --git a/store/go.mod b/store/go.mod
@@ -13,8 +13,8 @@ require (
 	github.com/cosmos/gogoproto v1.4.11
 	github.com/cosmos/iavl v1.0.0
 	github.com/cosmos/ics23/go v0.10.0
-	github.com/mattn/go-sqlite3 v1.14.17
 	github.com/linxGnu/grocksdb v1.8.5
+	github.com/mattn/go-sqlite3 v1.14.17
 	github.com/stretchr/testify v1.8.4
 	github.com/tidwall/btree v1.7.0
 	golang.org/x/exp v0.0.0-20231006140011-7918f672742d
@@ -56,13 +56,11 @@ require (
 	github.com/sasha-s/go-deadlock v0.3.1 // indirect
 	github.com/spf13/cast v1.5.1 // indirect
 	github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7 // indirect
-	golang.org/x/sync v0.4.0 // indirect
 	golang.org/x/crypto v0.15.0 // indirect
-	golang.org/x/mod v0.13.0 // indirect
 	golang.org/x/net v0.17.0 // indirect
+	golang.org/x/sync v0.4.0 // indirect
 	golang.org/x/sys v0.14.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
-	golang.org/x/tools v0.14.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
 	google.golang.org/grpc v1.59.0 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect

diff --git a/store/storage/sqlite/db.go b/store/storage/sqlite/db.go
@@ -20,6 +20,8 @@ const (
 	reservedStoreKey = "_RESERVED_"
 	keyLatestHeight  = "latest_height"
 
+	defaultBatchBufferSize = 100000
+
 	latestVersionStmt = `
 	INSERT INTO state_storage(store_key, key, value, version)
     VALUES(?, ?, ?, ?)
@@ -212,6 +214,40 @@ func (db *Database) ReverseIterator(storeKey string, version uint64, start, end
 	return newIterator(db.storage, storeKey, version, start, end, true)
 }
 
+// Restore implements the StorageSnapshotter interface.
+func (db *Database) Restore(version uint64, chStorage <-chan *store.KVPair) error {
+	latestVersion, err := db.GetLatestVersion()
+	if err != nil {
+		return fmt.Errorf("failed to get latest version: %w", err)
+	}
+	if version <= latestVersion {
+		return fmt.Errorf("the snapshot version %d is not greater than latest version %d", version, latestVersion)
+	}
+
+	b, err := NewBatch(db.storage, version)
+	if err != nil {
+		return err
+	}
+
+	for kvPair := range chStorage {
+		if err := b.Set(kvPair.StoreKey, kvPair.Key, kvPair.Value); err != nil {
+			return err
+		}
+
+		if b.Size() > defaultBatchBufferSize {
+			if err := b.Write(); err != nil {
+				return err
+			}
+		}
+	}
+
+	if err := b.Write(); err != nil {
+		return err
+	}
+
+	return db.SetLatestVersion(version)
+}
+
 func (db *Database) PrintRowsDebug() {
 	stmt, err := db.storage.Prepare("SELECT store_key, key, value, version, tombstone FROM state_storage")
 	if err != nil {