Add "buckets", "keys" and "get" CLI commands

These were previously removed but I find them quite useful so have
reintroduced them.

Signed-off-by: Lewis Marshall <lewis@lmars.net>

cmd/bolt/main.go

@@ -48,6 +48,18 @@ var (
 
 	// ErrPageIDRequired is returned when a required page id is not specified.
 	ErrPageIDRequired = errors.New("page id required")
+
+	// ErrBucketRequired is returned when a bucket is not specified.
+	ErrBucketRequired = errors.New("bucket required")
+
+	// ErrBucketNotFound is returned when a bucket is not found.
+	ErrBucketNotFound = errors.New("bucket not found")
+
+	// ErrKeyRequired is returned when a key is not specified.
+	ErrKeyRequired = errors.New("key required")
+
+	// ErrKeyNotFound is returned when a key is not found.
+	ErrKeyNotFound = errors.New("key not found")
 )
 
 // PageHeaderSize represents the size of the bolt.page header.
@@ -94,14 +106,20 @@ func (m *Main) Run(args ...string) error {
 		return ErrUsage
 	case "bench":
 		return newBenchCommand(m).Run(args[1:]...)
+	case "buckets":
+		return newBucketsCommand(m).Run(args[1:]...)
 	case "check":
 		return newCheckCommand(m).Run(args[1:]...)
 	case "compact":
 		return newCompactCommand(m).Run(args[1:]...)
 	case "dump":
 		return newDumpCommand(m).Run(args[1:]...)
+	case "get":
+		return newGetCommand(m).Run(args[1:]...)
 	case "info":
 		return newInfoCommand(m).Run(args[1:]...)
+	case "keys":
+		return newKeysCommand(m).Run(args[1:]...)
 	case "page":
 		return newPageCommand(m).Run(args[1:]...)
 	case "pages":
@@ -125,10 +143,13 @@ Usage:
 The commands are:
 
     bench       run synthetic benchmark against bolt
+    buckets     print a list of buckets
     check       verifies integrity of bolt database
     compact     copies a bolt database, compacting it in the process
     dump        print a hexidecimal dump of a single page
+    get         print the value of a key in a bucket
     info        print basic info
+    keys        print a list of keys in a bucket
     help        print this screen
     page        print one or more pages in human readable format
     pages       print list of pages with their types
@@ -867,6 +888,212 @@ experience corruption, please submit a ticket to the Bolt project page:
 `, "\n")
 }
 
+// BucketsCommand represents the "buckets" command execution.
+type BucketsCommand struct {
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewBucketsCommand returns a BucketsCommand.
+func newBucketsCommand(m *Main) *BucketsCommand {
+	return &BucketsCommand{
+		Stdin:  m.Stdin,
+		Stdout: m.Stdout,
+		Stderr: m.Stderr,
+	}
+}
+
+// Run executes the command.
+func (cmd *BucketsCommand) Run(args ...string) error {
+	// Parse flags.
+	fs := flag.NewFlagSet("", flag.ContinueOnError)
+	help := fs.Bool("h", false, "")
+	if err := fs.Parse(args); err != nil {
+		return err
+	} else if *help {
+		fmt.Fprintln(cmd.Stderr, cmd.Usage())
+		return ErrUsage
+	}
+
+	// Require database path.
+	path := fs.Arg(0)
+	if path == "" {
+		return ErrPathRequired
+	} else if _, err := os.Stat(path); os.IsNotExist(err) {
+		return ErrFileNotFound
+	}
+
+	// Open database.
+	db, err := bolt.Open(path, 0666, nil)
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	// Print buckets.
+	return db.View(func(tx *bolt.Tx) error {
+		return tx.ForEach(func(name []byte, _ *bolt.Bucket) error {
+			fmt.Fprintln(cmd.Stdout, string(name))
+			return nil
+		})
+	})
+}
+
+// Usage returns the help message.
+func (cmd *BucketsCommand) Usage() string {
+	return strings.TrimLeft(`
+usage: bolt buckets PATH
+
+Print a list of buckets.
+`, "\n")
+}
+
+// KeysCommand represents the "keys" command execution.
+type KeysCommand struct {
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewKeysCommand returns a KeysCommand.
+func newKeysCommand(m *Main) *KeysCommand {
+	return &KeysCommand{
+		Stdin:  m.Stdin,
+		Stdout: m.Stdout,
+		Stderr: m.Stderr,
+	}
+}
+
+// Run executes the command.
+func (cmd *KeysCommand) Run(args ...string) error {
+	// Parse flags.
+	fs := flag.NewFlagSet("", flag.ContinueOnError)
+	help := fs.Bool("h", false, "")
+	if err := fs.Parse(args); err != nil {
+		return err
+	} else if *help {
+		fmt.Fprintln(cmd.Stderr, cmd.Usage())
+		return ErrUsage
+	}
+
+	// Require database path and bucket.
+	path, bucket := fs.Arg(0), fs.Arg(1)
+	if path == "" {
+		return ErrPathRequired
+	} else if _, err := os.Stat(path); os.IsNotExist(err) {
+		return ErrFileNotFound
+	} else if bucket == "" {
+		return ErrBucketRequired
+	}
+
+	// Open database.
+	db, err := bolt.Open(path, 0666, nil)
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	// Print keys.
+	return db.View(func(tx *bolt.Tx) error {
+		// Find bucket.
+		b := tx.Bucket([]byte(bucket))
+		if b == nil {
+			return ErrBucketNotFound
+		}
+
+		// Iterate over each key.
+		return b.ForEach(func(key, _ []byte) error {
+			fmt.Fprintln(cmd.Stdout, string(key))
+			return nil
+		})
+	})
+}
+
+// Usage returns the help message.
+func (cmd *KeysCommand) Usage() string {
+	return strings.TrimLeft(`
+usage: bolt keys PATH BUCKET
+
+Print a list of keys in the given bucket.
+`, "\n")
+}
+
+// GetCommand represents the "get" command execution.
+type GetCommand struct {
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewGetCommand returns a GetCommand.
+func newGetCommand(m *Main) *GetCommand {
+	return &GetCommand{
+		Stdin:  m.Stdin,
+		Stdout: m.Stdout,
+		Stderr: m.Stderr,
+	}
+}
+
+// Run executes the command.
+func (cmd *GetCommand) Run(args ...string) error {
+	// Parse flags.
+	fs := flag.NewFlagSet("", flag.ContinueOnError)
+	help := fs.Bool("h", false, "")
+	if err := fs.Parse(args); err != nil {
+		return err
+	} else if *help {
+		fmt.Fprintln(cmd.Stderr, cmd.Usage())
+		return ErrUsage
+	}
+
+	// Require database path, bucket and key.
+	path, bucket, key := fs.Arg(0), fs.Arg(1), fs.Arg(2)
+	if path == "" {
+		return ErrPathRequired
+	} else if _, err := os.Stat(path); os.IsNotExist(err) {
+		return ErrFileNotFound
+	} else if bucket == "" {
+		return ErrBucketRequired
+	} else if key == "" {
+		return ErrKeyRequired
+	}
+
+	// Open database.
+	db, err := bolt.Open(path, 0666, nil)
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	// Print value.
+	return db.View(func(tx *bolt.Tx) error {
+		// Find bucket.
+		b := tx.Bucket([]byte(bucket))
+		if b == nil {
+			return ErrBucketNotFound
+		}
+
+		// Find value for given key.
+		val := b.Get([]byte(key))
+		if val == nil {
+			return ErrKeyNotFound
+		}
+
+		fmt.Fprintln(cmd.Stdout, string(val))
+		return nil
+	})
+}
+
+// Usage returns the help message.
+func (cmd *GetCommand) Usage() string {
+	return strings.TrimLeft(`
+usage: bolt get PATH BUCKET KEY
+
+Print the value of the given key in the given bucket.
+`, "\n")
+}
+
 var benchBucketName = []byte("bench")
 
 // BenchCommand represents the "bench" command execution.

cmd/bolt/main_test.go

@@ -146,6 +146,106 @@ func TestStatsCommand_Run(t *testing.T) {
 	}
 }
 
+// Ensure the "buckets" command can print a list of buckets.
+func TestBucketsCommand_Run(t *testing.T) {
+	db := MustOpen(0666, nil)
+	defer db.Close()
+
+	if err := db.Update(func(tx *bolt.Tx) error {
+		for _, name := range []string{"foo", "bar", "baz"} {
+			_, err := tx.CreateBucket([]byte(name))
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	}); err != nil {
+		t.Fatal(err)
+	}
+	db.DB.Close()
+
+	expected := "bar\nbaz\nfoo\n"
+
+	// Run the command.
+	m := NewMain()
+	if err := m.Run("buckets", db.Path); err != nil {
+		t.Fatal(err)
+	} else if actual := m.Stdout.String(); actual != expected {
+		t.Fatalf("unexpected stdout:\n\n%s", actual)
+	}
+}
+
+// Ensure the "keys" command can print a list of keys for a bucket.
+func TestKeysCommand_Run(t *testing.T) {
+	db := MustOpen(0666, nil)
+	defer db.Close()
+
+	if err := db.Update(func(tx *bolt.Tx) error {
+		for _, name := range []string{"foo", "bar"} {
+			b, err := tx.CreateBucket([]byte(name))
+			if err != nil {
+				return err
+			}
+			for i := 0; i < 3; i++ {
+				key := fmt.Sprintf("%s-%d", name, i)
+				if err := b.Put([]byte(key), []byte{0}); err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	}); err != nil {
+		t.Fatal(err)
+	}
+	db.DB.Close()
+
+	expected := "foo-0\nfoo-1\nfoo-2\n"
+
+	// Run the command.
+	m := NewMain()
+	if err := m.Run("keys", db.Path, "foo"); err != nil {
+		t.Fatal(err)
+	} else if actual := m.Stdout.String(); actual != expected {
+		t.Fatalf("unexpected stdout:\n\n%s", actual)
+	}
+}
+
+// Ensure the "get" command can print the value of a key in a bucket.
+func TestGetCommand_Run(t *testing.T) {
+	db := MustOpen(0666, nil)
+	defer db.Close()
+
+	if err := db.Update(func(tx *bolt.Tx) error {
+		for _, name := range []string{"foo", "bar"} {
+			b, err := tx.CreateBucket([]byte(name))
+			if err != nil {
+				return err
+			}
+			for i := 0; i < 3; i++ {
+				key := fmt.Sprintf("%s-%d", name, i)
+				val := fmt.Sprintf("val-%s-%d", name, i)
+				if err := b.Put([]byte(key), []byte(val)); err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	}); err != nil {
+		t.Fatal(err)
+	}
+	db.DB.Close()
+
+	expected := "val-foo-1\n"
+
+	// Run the command.
+	m := NewMain()
+	if err := m.Run("get", db.Path, "foo", "foo-1"); err != nil {
+		t.Fatal(err)
+	} else if actual := m.Stdout.String(); actual != expected {
+		t.Fatalf("unexpected stdout:\n\n%s", actual)
+	}
+}
+
 // Main represents a test wrapper for main.Main that records output.
 type Main struct {
 	*main.Main