Skip to content

Commit

Permalink
cli: add command for traversing MPT for the current state
Browse files Browse the repository at this point in the history
Traverse MPT for the current state and dump key/value pairs into file.

Close #3519

Signed-off-by: Ekaterina Pavlova <ekt@morphbits.io>
  • Loading branch information
AliceInHunterland committed Jul 24, 2024
1 parent f0ae14e commit 4c847ec
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 0 deletions.
98 changes: 98 additions & 0 deletions cli/server/mptdump.go
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}

Check warning on line 27 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L24-L27

Added lines #L24 - L27 were not covered by tests

store.Seek(rng, func(k, v []byte) bool {
kvPairs = append(kvPairs, KVPair{
Key: hex.EncodeToString(k),
Value: hex.EncodeToString(v),
})
return true
})

Check warning on line 35 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L29-L35

Added lines #L29 - L35 were not covered by tests

return kvPairs, nil

Check warning on line 37 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L37

Added line #L37 was not covered by tests
}

// 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

Check warning on line 44 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L41-L44

Added lines #L41 - L44 were not covered by tests
}
defer file.Close()

Check warning on line 46 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L46

Added line #L46 was not covered by tests

encoder := json.NewEncoder(file)
encoder.SetIndent("", " ")
err = encoder.Encode(kvPairs)
if err != nil {
return err

Check warning on line 52 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L48-L52

Added lines #L48 - L52 were not covered by tests
}

return nil

Check warning on line 55 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L55

Added line #L55 was not covered by tests
}

// 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)

Check warning on line 63 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L59-L63

Added lines #L59 - L63 were not covered by tests
}

chain, store, err := initBlockChain(cfg, logger)
if err != nil {
return cli.Exit(err, 1)

Check warning on line 68 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L66-L68

Added lines #L66 - L68 were not covered by tests
}
defer store.Close()
defer chain.Close()

Check warning on line 71 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L70-L71

Added lines #L70 - L71 were not covered by tests

stateModule := chain.GetStateModule()
stateRoot := stateModule.CurrentLocalStateRoot()
stateRootHash := stateRoot

Check warning on line 75 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L73-L75

Added lines #L73 - L75 were not covered by tests

trieStore := mpt.NewTrieStore(stateRootHash, mpt.ModeLatest, store)

Check warning on line 77 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L77

Added line #L77 was not covered by tests

startTime := time.Now()
kvPairs, err := TraverseMPT(trieStore)
if err != nil {
return cli.Exit(err, 1)

Check warning on line 82 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L79-L82

Added lines #L79 - L82 were not covered by tests
}

outputFile := ctx.String("out")
if outputFile == "" {
outputFile = "kv_pairs.json"

Check warning on line 87 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L85-L87

Added lines #L85 - L87 were not covered by tests
}

err = DumpKVPairsToFile(kvPairs, outputFile)
if err != nil {
return cli.Exit(err, 1)

Check warning on line 92 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L90-L92

Added lines #L90 - L92 were not covered by tests
}

duration := time.Since(startTime)
fmt.Printf("MPT key-value pairs successfully dumped to %s in %s\n", outputFile, duration)
return nil

Check warning on line 97 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L95-L97

Added lines #L95 - L97 were not covered by tests
}
13 changes: 13 additions & 0 deletions cli/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,19 @@ func NewCommands() []*cli.Command {
Action: resetDB,
Flags: cfgHeightFlags,
},
{
Name: "traverse",
Usage: "Traverse the MPT and dump key-value pairs to a file",
UsageText: "neo-go db traverse [--out file] [--config-path path] [-p/-m/-t] [--config-file file]",
Action: traverseMPT,
Flags: append(cfgFlags,
&cli.StringFlag{
Name: "out",
Aliases: []string{"o"},
Usage: "Output file (default: kv_pairs.json)",
},
),
},
},
},
}
Expand Down

0 comments on commit 4c847ec

Please sign in to comment.