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.
Example usage: `./bin/neo-go db traverse --config-file config/protocol
.mainnet.yml`.

Close #3519

Signed-off-by: Ekaterina Pavlova <ekt@morphbits.io>
  • Loading branch information
AliceInHunterland committed Jul 29, 2024
1 parent 8bc06ec commit ea10fc6
Show file tree
Hide file tree
Showing 2 changed files with 100 additions and 0 deletions.
87 changes: 87 additions & 0 deletions cli/server/mptdump.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
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, encoder *json.Encoder) error {
prefix := []byte{byte(storage.STStorage)}
rng := storage.SeekRange{Prefix: prefix}

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L24-L26

Added lines #L24 - L26 were not covered by tests

store.Seek(rng, func(k, v []byte) bool {
kvPair := KVPair{
Key: hex.EncodeToString(k),
Value: hex.EncodeToString(v),

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L28-L31

Added lines #L28 - L31 were not covered by tests
}
if err := encoder.Encode(kvPair); err != nil {
fmt.Printf("error encoding key-value pair: %v\n", err)
return false

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L33-L35

Added lines #L33 - L35 were not covered by tests
}
return true

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
})
return nil

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L39

Added line #L39 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 47 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L43-L47

Added lines #L43 - L47 were not covered by tests
}

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

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L50-L52

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

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L54-L55

Added lines #L54 - L55 were not covered by tests

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

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L57-L59

Added lines #L57 - L59 were not covered by tests

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

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L61

Added line #L61 was not covered by tests

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

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L63-L65

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

file, err := os.Create(outputFile)
if err != nil {
return fmt.Errorf("error creating file: %w", err)

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L68-L70

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

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L72

Added line #L72 was not covered by tests

encoder := json.NewEncoder(file)
encoder.SetIndent("", " ")

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L74-L75

Added lines #L74 - L75 were not covered by tests

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

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

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L77-L81

Added lines #L77 - L81 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 86 in cli/server/mptdump.go

View check run for this annotation

Codecov / codecov/patch

cli/server/mptdump.go#L84-L86

Added lines #L84 - L86 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 ea10fc6

Please sign in to comment.