Add dbtool

Author: kolesnikovae

go.mod

@@ -45,6 +45,7 @@ require (
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/shirou/gopsutil v3.21.4+incompatible
 	github.com/sirupsen/logrus v1.7.0
+	github.com/spf13/cobra v0.0.5
 	github.com/tklauser/go-sysconf v0.3.6 // indirect
 	github.com/twmb/murmur3 v1.1.5
 	github.com/valyala/bytebufferpool v1.0.0

go.sum

@@ -205,6 +205,7 @@ github.com/ianlancetaylor/demangle v0.0.0-20200715173712-053cf528c12f/go.mod h1:
 github.com/imdario/mergo v0.3.8/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
 github.com/imdario/mergo v0.3.11 h1:3tnifQM4i+fbajXKBHXWEH+KvNHqojZ778UH75j3bGA=
 github.com/imdario/mergo v0.3.11/go.mod h1:jmQim1M+e3UYxmgPu/WyfjB3N3VflVyUjjjwH0dnCYA=
+github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/joho/godotenv v1.3.0 h1:Zjp+RcGpHhGlrMbJzXTrZZPrWj+1vfm90La1wgB6Bhc=
 github.com/joho/godotenv v1.3.0/go.mod h1:7hK45KPybAkOC6peb+G5yklZfMxEjkZhHbwpqxOKXbg=
@@ -341,8 +342,10 @@ github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0b
 github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/spf13/afero v1.1.2/go.mod h1:j4pytiNVoe2o6bmDsKpLACNPDBIoEAkihy7loJ1B0CQ=
 github.com/spf13/cast v1.3.0/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
+github.com/spf13/cobra v0.0.5 h1:f0B+LkLX6DtmRH1isoNA9VTtNUK9K8xYd28JNNfOv/s=
 github.com/spf13/cobra v0.0.5/go.mod h1:3K3wKZymM7VvHMDS9+Akkh4K60UwM26emMESw8tLCHU=
 github.com/spf13/jwalterweatherman v1.0.0/go.mod h1:cQK4TGJAtQXfYWX+Ddv3mKDzgVb68N+wFjFa4jdeBTo=
+github.com/spf13/pflag v1.0.3 h1:zPAT6CGy6wXeQ7NtTnaTerfKOsV6V6F8agHXFiazDkg=
 github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
 github.com/spf13/viper v1.3.2/go.mod h1:ZiWeW+zYFKm7srdB9IoDzzZXaJaI5eL9QjNiN/DMA2s=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

scripts/dbtool/cmd/get.go

@@ -0,0 +1,35 @@
+package cmd
+
+import (
+	"bytes"
+	"io"
+	"os"
+
+	"github.com/dgraph-io/badger/v2"
+	"github.com/spf13/cobra"
+)
+
+func (d *dbTool) newGetCommand() *cobra.Command {
+	return &cobra.Command{
+		Use:     "get <key>",
+		Short:   "Retrieves key value and dumps it to stdout",
+		RunE:    d.runGet,
+		Args:    cobra.MinimumNArgs(1),
+		PreRunE: d.openDB(true),
+	}
+}
+
+func (d *dbTool) runGet(_ *cobra.Command, args []string) error {
+	return d.db.View(func(txn *badger.Txn) error {
+		item, err := txn.Get([]byte(args[0]))
+		if err != nil {
+			return err
+		}
+		v, err := item.ValueCopy(nil)
+		if err != nil {
+			return err
+		}
+		_, _ = io.Copy(os.Stdout, bytes.NewBuffer(v))
+		return nil
+	})
+}

scripts/dbtool/cmd/list.go

@@ -0,0 +1,34 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/dgraph-io/badger/v2"
+	"github.com/spf13/cobra"
+)
+
+func (d *dbTool) newListCommand() *cobra.Command {
+	cmd := cobra.Command{
+		Use:     "list",
+		Short:   "Lists db keys",
+		RunE:    d.runList,
+		PreRunE: d.openDB(true),
+	}
+	cmd.Flags().StringVarP(&d.prefix, "prefix", "p", "", "key prefix")
+	return &cmd
+}
+
+func (d *dbTool) runList(_ *cobra.Command, _ []string) error {
+	return d.db.View(func(txn *badger.Txn) error {
+		opts := badger.DefaultIteratorOptions
+		opts.PrefetchValues = false
+		opts.Prefix = []byte(d.prefix)
+		it := txn.NewIterator(opts)
+		defer it.Close()
+		for it.Rewind(); it.Valid(); it.Next() {
+			_, _ = fmt.Fprintf(os.Stdout, "%s\n", it.Item().Key())
+		}
+		return nil
+	})
+}

scripts/dbtool/cmd/remove.go

@@ -0,0 +1,22 @@
+package cmd
+
+import (
+	"github.com/dgraph-io/badger/v2"
+	"github.com/spf13/cobra"
+)
+
+func (d *dbTool) newRemoveCommand() *cobra.Command {
+	return &cobra.Command{
+		Use:     "remove <key>",
+		Short:   "Removes key value from the database",
+		RunE:    d.runRemove,
+		Args:    cobra.MinimumNArgs(1),
+		PreRunE: d.openDB(false),
+	}
+}
+
+func (d *dbTool) runRemove(_ *cobra.Command, args []string) error {
+	return d.db.Update(func(txn *badger.Txn) error {
+		return txn.Delete([]byte(args[0]))
+	})
+}

scripts/dbtool/cmd/root.go

@@ -0,0 +1,80 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/dgraph-io/badger/v2"
+	"github.com/dgraph-io/badger/v2/options"
+	"github.com/spf13/cobra"
+)
+
+type dbTool struct {
+	dir    string
+	prefix string
+
+	db *badger.DB
+}
+
+func Execute() error {
+	return newRootCommand().Execute()
+}
+
+func newRootCommand() *cobra.Command {
+	d := new(dbTool)
+	root := cobra.Command{
+		Use:                "dbtool [command]",
+		Short:              "DB helper tool",
+		SilenceUsage:       true,
+		SilenceErrors:      true,
+		PersistentPostRunE: d.closeDB(),
+	}
+
+	root.PersistentFlags().StringVarP(&d.dir, "dir", "d", ".", "database directory path")
+	root.AddCommand(
+		d.newListCommand(),
+		d.newRemoveCommand(),
+		d.newGetCommand())
+
+	return &root
+}
+
+func (d *dbTool) openDB(ro bool) func(cmd *cobra.Command, args []string) error {
+	return func(c *cobra.Command, _ []string) error {
+		db, err := openDB(d.dir, ro)
+		if err != nil {
+			return fmt.Errorf("failed to open %s: %v", d.dir, err)
+		}
+		d.db = db
+		return nil
+	}
+}
+
+func (d *dbTool) closeDB() func(cmd *cobra.Command, args []string) error {
+	return func(c *cobra.Command, _ []string) error {
+		if d.db == nil {
+			return nil
+		}
+		if err := d.db.Close(); err != nil {
+			return fmt.Errorf("closing database: %w", err)
+		}
+		return nil
+	}
+}
+
+func openDB(dir string, ro bool) (*badger.DB, error) {
+	if !ro {
+		// When DB is opened for RW, badger.Open does not fail
+		// if the directory is not a badger database, instead it
+		// creates one.
+		f, err := os.Open(filepath.Join(dir, badger.ManifestFilename))
+		if err != nil {
+			return nil, err
+		}
+		_ = f.Close()
+	}
+	return badger.Open(badger.DefaultOptions(dir).
+		WithCompression(options.ZSTD).
+		WithReadOnly(ro))
+}

scripts/dbtool/main.go

@@ -0,0 +1,17 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/fatih/color"
+
+	"github.com/pyroscope-io/pyroscope/scripts/dbtool/cmd"
+)
+
+func main() {
+	if err := cmd.Execute(); err != nil {
+		_, _ = fmt.Fprintln(os.Stderr, color.RedString("Error:"), err)
+		os.Exit(1)
+	}
+}