tool to search db for a key prefix (#11417)

* tool to query db by a key prefix

* cleanup

* lint and fmt

* db/kv public visibility

We've discussed before that Bazel visibility constraints don't
accomplish much in go, so I'm phasing them out in places where they get
in the way.

* derp

Co-authored-by: Kasey Kirkham <kasey@users.noreply.github.com>
Co-authored-by: Radosław Kapka <rkapka@wp.pl>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

cmd/prysmctl/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
     visibility = ["//visibility:private"],
     deps = [
         "//cmd/prysmctl/checkpointsync:go_default_library",
+        "//cmd/prysmctl/db:go_default_library",
         "//cmd/prysmctl/deprecated:go_default_library",
         "//cmd/prysmctl/p2p:go_default_library",
         "//cmd/prysmctl/signing:go_default_library",

cmd/prysmctl/db/BUILD.bazel

@@ -0,0 +1,21 @@
+load("@prysm//tools/go:def.bzl", "go_library")
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "buckets.go",
+        "cmd.go",
+        "query.go",
+    ],
+    importpath = "github.com/prysmaticlabs/prysm/v3/cmd/prysmctl/db",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//beacon-chain/db/kv:go_default_library",
+        "//config/params:go_default_library",
+        "@com_github_ethereum_go_ethereum//common/hexutil:go_default_library",
+        "@com_github_pkg_errors//:go_default_library",
+        "@com_github_sirupsen_logrus//:go_default_library",
+        "@com_github_urfave_cli_v2//:go_default_library",
+        "@io_etcd_go_bbolt//:go_default_library",
+    ],
+)

cmd/prysmctl/db/buckets.go

@@ -0,0 +1,32 @@
+package db
+
+import (
+	"fmt"
+
+	"github.com/prysmaticlabs/prysm/v3/beacon-chain/db/kv"
+	"github.com/urfave/cli/v2"
+)
+
+var bucketsFlags = struct {
+	Path string
+}{}
+
+var bucketsCmd = &cli.Command{
+	Name:   "buckets",
+	Usage:  "list db buckets",
+	Action: bucketsAction,
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:        "path",
+			Usage:       "path to directory containing beaconchain.db",
+			Destination: &bucketsFlags.Path,
+		},
+	},
+}
+
+func bucketsAction(_ *cli.Context) error {
+	for _, b := range kv.Buckets {
+		fmt.Printf("%s\n", string(b))
+	}
+	return nil
+}

cmd/prysmctl/db/cmd.go

@@ -0,0 +1,14 @@
+package db
+
+import "github.com/urfave/cli/v2"
+
+var Commands = []*cli.Command{
+	{
+		Name:  "db",
+		Usage: "commands to work with the prysm beacon db",
+		Subcommands: []*cli.Command{
+			queryCmd,
+			bucketsCmd,
+		},
+	},
+}

cmd/prysmctl/db/query.go

@@ -0,0 +1,97 @@
+package db
+
+import (
+	"bytes"
+	"fmt"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common/hexutil"
+	"github.com/pkg/errors"
+	"github.com/prysmaticlabs/prysm/v3/config/params"
+	log "github.com/sirupsen/logrus"
+	"github.com/urfave/cli/v2"
+	bolt "go.etcd.io/bbolt"
+)
+
+var queryFlags = struct {
+	Path     string
+	Bucket   string
+	KeysOnly bool
+	Prefix   string
+}{}
+
+var queryCmd = &cli.Command{
+	Name:   "query",
+	Usage:  "database query tool",
+	Action: queryAction,
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:        "bucket",
+			Usage:       "boltdb bucket to search",
+			Destination: &queryFlags.Bucket,
+		},
+		&cli.StringFlag{
+			Name:        "path",
+			Usage:       "path to directory containing beaconchain.db",
+			Destination: &queryFlags.Path,
+		},
+		&cli.StringFlag{
+			Name:        "prefix",
+			Usage:       "prefix of db record key to match against (eg 0xa1 would match 0xa10, 0xa1f etc)",
+			Destination: &queryFlags.Prefix,
+		},
+		&cli.BoolFlag{
+			Name:        "print-keys",
+			Usage:       "only display keys, not values",
+			Destination: &queryFlags.KeysOnly,
+		},
+	},
+}
+
+func queryAction(_ *cli.Context) error {
+	flags := queryFlags
+	db, err := getDB(flags.Path)
+	if err != nil {
+		return err
+	}
+	if flags.Prefix != "" {
+		return prefixScan(db, flags.Bucket, flags.Prefix, flags.KeysOnly)
+	}
+	return nil
+}
+
+func prefixScan(db *bolt.DB, bucket, prefix string, keysOnly bool) error {
+	if !keysOnly {
+		return errors.New("prefix scan with value display not implemented")
+	}
+	pb, err := hexutil.Decode(prefix)
+	if err != nil {
+		return err
+	}
+	log.Infof("scanning for prefix=%#x", pb)
+	return db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(bucket))
+		c := b.Cursor()
+		for k, _ := c.Seek(pb); k != nil && bytes.HasPrefix(k, pb); k, _ = c.Next() {
+			fmt.Printf("%#x\n", k)
+		}
+		return nil
+	})
+}
+
+func getDB(path string) (*bolt.DB, error) {
+	bdb, err := bolt.Open(
+		path,
+		params.BeaconIoConfig().ReadWritePermissions,
+		&bolt.Options{
+			Timeout: 1 * time.Second,
+		},
+	)
+	if err != nil {
+		if errors.Is(err, bolt.ErrTimeout) {
+			return nil, errors.New("cannot obtain database lock, database may be in use by another process")
+		}
+		return nil, err
+	}
+	return bdb, nil
+}

cmd/prysmctl/main.go

@@ -4,6 +4,7 @@ import (
 	"os"
 
 	"github.com/prysmaticlabs/prysm/v3/cmd/prysmctl/checkpointsync"
+	"github.com/prysmaticlabs/prysm/v3/cmd/prysmctl/db"
 	"github.com/prysmaticlabs/prysm/v3/cmd/prysmctl/deprecated"
 	"github.com/prysmaticlabs/prysm/v3/cmd/prysmctl/p2p"
 	"github.com/prysmaticlabs/prysm/v3/cmd/prysmctl/signing"
@@ -31,6 +32,7 @@ func init() {
 	prysmctlCommands = append(prysmctlCommands, deprecated.Commands...)
 
 	prysmctlCommands = append(prysmctlCommands, checkpointsync.Commands...)
+	prysmctlCommands = append(prysmctlCommands, db.Commands...)
 	prysmctlCommands = append(prysmctlCommands, p2p.Commands...)
 	prysmctlCommands = append(prysmctlCommands, testnet.Commands...)
 	prysmctlCommands = append(prysmctlCommands, weaksubjectivity.Commands...)