github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/scripts/dbtool/cmd/root.go (about)

     1  package cmd
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"path/filepath"
     7  
     8  	"github.com/dgraph-io/badger/v2"
     9  	"github.com/dgraph-io/badger/v2/options"
    10  	"github.com/spf13/cobra"
    11  )
    12  
    13  type dbTool struct {
    14  	dir    string
    15  	prefix string
    16  
    17  	db *badger.DB
    18  }
    19  
    20  func Execute() error {
    21  	return newRootCommand().Execute()
    22  }
    23  
    24  func newRootCommand() *cobra.Command {
    25  	d := new(dbTool)
    26  	root := cobra.Command{
    27  		Use:                "dbtool [command]",
    28  		Short:              "DB helper tool",
    29  		SilenceUsage:       true,
    30  		SilenceErrors:      true,
    31  		PersistentPostRunE: d.closeDB(),
    32  	}
    33  
    34  	root.PersistentFlags().StringVarP(&d.dir, "dir", "d", ".", "database directory path")
    35  	root.AddCommand(
    36  		d.newListCommand(),
    37  		d.newRemoveCommand(),
    38  		d.newGetCommand())
    39  
    40  	return &root
    41  }
    42  
    43  func (d *dbTool) openDB(ro bool) func(cmd *cobra.Command, args []string) error {
    44  	return func(c *cobra.Command, _ []string) error {
    45  		db, err := openDB(d.dir, ro)
    46  		if err != nil {
    47  			return fmt.Errorf("failed to open %s: %v", d.dir, err)
    48  		}
    49  		d.db = db
    50  		return nil
    51  	}
    52  }
    53  
    54  func (d *dbTool) closeDB() func(cmd *cobra.Command, args []string) error {
    55  	return func(c *cobra.Command, _ []string) error {
    56  		if d.db == nil {
    57  			return nil
    58  		}
    59  		if err := d.db.Close(); err != nil {
    60  			return fmt.Errorf("closing database: %w", err)
    61  		}
    62  		return nil
    63  	}
    64  }
    65  
    66  func openDB(dir string, ro bool) (*badger.DB, error) {
    67  	if !ro {
    68  		// When DB is opened for RW, badger.Open does not fail
    69  		// if the directory is not a badger database, instead it
    70  		// creates one.
    71  		f, err := os.Open(filepath.Join(dir, badger.ManifestFilename))
    72  		if err != nil {
    73  			return nil, err
    74  		}
    75  		_ = f.Close()
    76  	}
    77  	return badger.Open(badger.DefaultOptions(dir).
    78  		WithCompression(options.ZSTD).
    79  		WithReadOnly(ro))
    80  }