-
Notifications
You must be signed in to change notification settings - Fork 81
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
cli: add command for traversing MPT for the current state
Traverse MPT for the current state and dump key/value pairs into file. Close #3519 Signed-off-by: Ekaterina Pavlova <[email protected]>
- Loading branch information
1 parent
f0ae14e
commit 4c847ec
Showing
2 changed files
with
111 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,98 @@ | ||
package server | ||
|
||
import ( | ||
"encoding/hex" | ||
"encoding/json" | ||
"fmt" | ||
"os" | ||
"time" | ||
|
||
"github.com/nspcc-dev/neo-go/cli/options" | ||
"github.com/nspcc-dev/neo-go/pkg/core/mpt" | ||
"github.com/nspcc-dev/neo-go/pkg/core/storage" | ||
"github.com/urfave/cli/v2" | ||
"go.uber.org/zap" | ||
) | ||
|
||
// KVPair represents a key-value pair. | ||
type KVPair struct { | ||
Key string `json:"key"` | ||
Value string `json:"value"` | ||
} | ||
|
||
// TraverseMPT collects key-value pairs from the TrieStore and returns them. | ||
func TraverseMPT(store *mpt.TrieStore) ([]KVPair, error) { | ||
var kvPairs []KVPair | ||
prefix := []byte{byte(storage.STStorage)} | ||
rng := storage.SeekRange{Prefix: prefix} | ||
|
||
store.Seek(rng, func(k, v []byte) bool { | ||
kvPairs = append(kvPairs, KVPair{ | ||
Key: hex.EncodeToString(k), | ||
Value: hex.EncodeToString(v), | ||
}) | ||
return true | ||
}) | ||
|
||
return kvPairs, nil | ||
} | ||
|
||
// DumpKVPairsToFile dumps the collected key-value pairs into a file. | ||
func DumpKVPairsToFile(kvPairs []KVPair, filename string) error { | ||
file, err := os.Create(filename) | ||
if err != nil { | ||
return err | ||
} | ||
defer file.Close() | ||
|
||
encoder := json.NewEncoder(file) | ||
encoder.SetIndent("", " ") | ||
err = encoder.Encode(kvPairs) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
return nil | ||
} | ||
|
||
// traverseMPT handles the CLI command to traverse the MPT and dump key-value pairs. | ||
func traverseMPT(ctx *cli.Context) error { | ||
logger := zap.NewExample() | ||
cfg, err := options.GetConfigFromContext(ctx) | ||
if err != nil { | ||
return cli.Exit(err, 1) | ||
} | ||
|
||
chain, store, err := initBlockChain(cfg, logger) | ||
if err != nil { | ||
return cli.Exit(err, 1) | ||
} | ||
defer store.Close() | ||
defer chain.Close() | ||
|
||
stateModule := chain.GetStateModule() | ||
stateRoot := stateModule.CurrentLocalStateRoot() | ||
stateRootHash := stateRoot | ||
|
||
trieStore := mpt.NewTrieStore(stateRootHash, mpt.ModeLatest, store) | ||
|
||
startTime := time.Now() | ||
kvPairs, err := TraverseMPT(trieStore) | ||
if err != nil { | ||
return cli.Exit(err, 1) | ||
} | ||
|
||
outputFile := ctx.String("out") | ||
if outputFile == "" { | ||
outputFile = "kv_pairs.json" | ||
} | ||
|
||
err = DumpKVPairsToFile(kvPairs, outputFile) | ||
if err != nil { | ||
return cli.Exit(err, 1) | ||
} | ||
|
||
duration := time.Since(startTime) | ||
fmt.Printf("MPT key-value pairs successfully dumped to %s in %s\n", outputFile, duration) | ||
return nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters