aggregate bucket stats recursively and add stats to cmd

3 files changed, 55 insertions, 12 deletions

diff --git a/cmd/bolt/main.go b/cmd/bolt/main.go
index 659f1c3..302fe09 100644
--- a/cmd/bolt/main.go
+++ b/cmd/bolt/main.go
@@ -93,7 +93,7 @@ func NewApp() *cli.App {
 		},
 		{
 			Name:  "stats",
-			Usage: "Retrieve statistics for a bucket (aggregated recursively)",
+			Usage: "Aggregate statistics for all buckets matching specified prefix",
 			Action: func(c *cli.Context) {
 				path, name := c.Args().Get(0), c.Args().Get(1)
 				Stats(path, name)
diff --git a/cmd/bolt/stats.go b/cmd/bolt/stats.go
index da344d0..6f8505a 100644
--- a/cmd/bolt/stats.go
+++ b/cmd/bolt/stats.go
@@ -1,13 +1,14 @@
 package main
 
 import (
+	"bytes"
 	"os"
 
 	"github.com/boltdb/bolt"
 )
 
-// Keys retrieves a list of keys for a given bucket.
-func Stats(path, name string) {
+// Collect stats for all top level buckets matching the prefix.
+func Stats(path, prefix string) {
 	if _, err := os.Stat(path); os.IsNotExist(err) {
 		fatal(err)
 		return
@@ -21,15 +22,18 @@ func Stats(path, name string) {
 	defer db.Close()
 
 	err = db.View(func(tx *bolt.Tx) error {
-		// Find bucket.
-		b := tx.Bucket([]byte(name))
-		if b == nil {
-			fatalf("bucket not found: %s", name)
+		var s bolt.BucketStats
+		var count int
+		var prefix = []byte(prefix)
+		tx.ForEach(func(name []byte, b *bolt.Bucket) error {
+			if bytes.HasPrefix(name, prefix) {
+				s.Add(b.Stats())
+				count += 1
+			}
 			return nil
-		}
+		})
+		printf("Aggregate statistics for %d buckets\n\n", count)
 
-		// Iterate over each key.
-		s := b.Stats()
 		println("Page count statistics")
 		printf("\tNumber of logical branch pages: %d\n", s.BranchPageN)
 		printf("\tNumber of physical branch overflow pages: %d\n", s.BranchOverflowN)
@@ -42,9 +46,11 @@ func Stats(path, name string) {
 
 		println("Page size utilization")
 		printf("\tBytes allocated for physical branch pages: %d\n", s.BranchAlloc)
-		printf("\tBytes actually used for branch data: %d\n", s.BranchInuse)
+		percentage := int(float32(s.BranchInuse) * 100.0 / float32(s.BranchAlloc))
+		printf("\tBytes actually used for branch data: %d (%d%%)\n", s.BranchInuse, percentage)
 		printf("\tBytes allocated for physical leaf pages: %d\n", s.LeafAlloc)
-		printf("\tBytes actually used for leaf data: %d\n", s.LeafInuse)
+		percentage = int(float32(s.LeafInuse) * 100.0 / float32(s.LeafAlloc))
+		printf("\tBytes actually used for leaf data: %d (%d%%)\n", s.LeafInuse, percentage)
 		return nil
 	})
 	if err != nil {
diff --git a/cmd/bolt/stats_test.go b/cmd/bolt/stats_test.go
new file mode 100644
index 0000000..9662f09
--- /dev/null
+++ b/cmd/bolt/stats_test.go
@@ -0,0 +1,37 @@
+package main_test
+
+import (
+	"testing"
+
+	"github.com/boltdb/bolt"
+	. "github.com/boltdb/bolt/cmd/bolt"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestStats(t *testing.T) {
+	SetTestMode(true)
+	open(func(db *bolt.DB, path string) {
+		db.Update(func(tx *bolt.Tx) error {
+			tx.CreateBucket([]byte("foo"))
+			tx.CreateBucket([]byte("bar"))
+			tx.CreateBucket([]byte("baz"))
+			return nil
+		})
+		db.Close()
+		output := run("stats", path, "b")
+		assert.Equal(t, "Aggregate statistics for 2 buckets\n\n"+
+			"Page count statistics\n"+
+			"\tNumber of logical branch pages: 0\n"+
+			"\tNumber of physical branch overflow pages: 0\n"+
+			"\tNumber of logical leaf pages: 2\n"+
+			"\tNumber of physical leaf overflow pages: 0\n"+
+			"Tree statistics\n"+
+			"\tNumber of keys/value pairs: 0\n"+
+			"\tNumber of levels in B+tree: 0\n"+
+			"Page size utilization\n"+
+			"\tBytes allocated for physical branch pages: 0\n"+
+			"\tBytes actually used for branch data: 0\n"+
+			"\tBytes allocated for physical leaf pages: 8192\n"+
+			"\tBytes actually used for leaf data: 0", output)
+	})
+}