add method Inspect to inspect bucket structure

Also added a related command: bbolt inspect db

The outputed etcd data structure:
{
    "name": "root",
    "keyN": 0,
    "children": [
        {
            "name": "alarm",
            "keyN": 0
        },
        {
            "name": "auth",
            "keyN": 2
        },
        {
            "name": "authRoles",
            "keyN": 1
        },
        {
            "name": "authUsers",
            "keyN": 1
        },
        {
            "name": "cluster",
            "keyN": 1
        },
        {
            "name": "key",
            "keyN": 1285
        },
        {
            "name": "lease",
            "keyN": 2
        },
        {
            "name": "members",
            "keyN": 1
        },
        {
            "name": "members_removed",
            "keyN": 0
        },
        {
            "name": "meta",
            "keyN": 3
        }
    ]
}

Signed-off-by: Benjamin Wang <[email protected]>

Author: ahrtr

bucket.go

@@ -392,6 +392,30 @@ func (b *Bucket) MoveBucket(key []byte, dstBucket *Bucket) (err error) {
 	return nil
 }
 
+// Inspect returns the structure of the bucket.
+func (b *Bucket) Inspect() BucketStructure {
+	return b.recursivelyInspect([]byte("root"))
+}
+
+func (b *Bucket) recursivelyInspect(name []byte) BucketStructure {
+	bs := BucketStructure{Name: string(name)}
+
+	keyN := 0
+	c := b.Cursor()
+	for k, _, flags := c.first(); k != nil; k, _, flags = c.next() {
+		if flags&common.BucketLeafFlag != 0 {
+			childBucket := b.Bucket(k)
+			childBS := childBucket.recursivelyInspect(k)
+			bs.Children = append(bs.Children, childBS)
+		} else {
+			keyN++
+		}
+	}
+	bs.KeyN = keyN
+
+	return bs
+}
+
 // Get retrieves the value for a key in the bucket.
 // Returns a nil value if the key does not exist or if the key is a nested bucket.
 // The returned value is only valid for the life of the transaction.
@@ -955,3 +979,9 @@ func cloneBytes(v []byte) []byte {
 	copy(clone, v)
 	return clone
 }
+
+type BucketStructure struct {
+	Name     string            `json:"name"`              // name of the bucket
+	KeyN     int               `json:"keyN"`              // number of key/value pairs
+	Children []BucketStructure `json:"buckets,omitempty"` // child buckets
+}

bucket_test.go

@@ -1623,6 +1623,111 @@ func TestBucket_Stats_Nested(t *testing.T) {
 	}
 }
 
+func TestBucket_Inspect(t *testing.T) {
+	db := btesting.MustCreateDB(t)
+
+	expectedStructure := bolt.BucketStructure{
+		Name: "root",
+		KeyN: 0,
+		Children: []bolt.BucketStructure{
+			{
+				Name: "b1",
+				KeyN: 3,
+				Children: []bolt.BucketStructure{
+					{
+						Name: "b1_1",
+						KeyN: 6,
+					},
+					{
+						Name: "b1_2",
+						KeyN: 7,
+					},
+					{
+						Name: "b1_3",
+						KeyN: 8,
+					},
+				},
+			},
+			{
+				Name: "b2",
+				KeyN: 4,
+				Children: []bolt.BucketStructure{
+					{
+						Name: "b2_1",
+						KeyN: 10,
+					},
+					{
+						Name: "b2_2",
+						KeyN: 12,
+						Children: []bolt.BucketStructure{
+							{
+								Name: "b2_2_1",
+								KeyN: 2,
+							},
+							{
+								Name: "b2_2_2",
+								KeyN: 3,
+							},
+						},
+					},
+					{
+						Name: "b2_3",
+						KeyN: 11,
+					},
+				},
+			},
+		},
+	}
+
+	type bucketItem struct {
+		b  *bolt.Bucket
+		bs bolt.BucketStructure
+	}
+
+	t.Log("Populating the database")
+	err := db.Update(func(tx *bolt.Tx) error {
+		queue := []bucketItem{
+			{
+				b:  nil,
+				bs: expectedStructure,
+			},
+		}
+
+		for len(queue) > 0 {
+			item := queue[0]
+			queue = queue[1:]
+
+			if item.b != nil {
+				for i := 0; i < item.bs.KeyN; i++ {
+					err := item.b.Put([]byte(fmt.Sprintf("%02d", i)), []byte(fmt.Sprintf("%02d", i)))
+					require.NoError(t, err)
+				}
+
+				for _, child := range item.bs.Children {
+					childBucket, err := item.b.CreateBucket([]byte(child.Name))
+					require.NoError(t, err)
+					queue = append(queue, bucketItem{b: childBucket, bs: child})
+				}
+			} else {
+				for _, child := range item.bs.Children {
+					childBucket, err := tx.CreateBucket([]byte(child.Name))
+					require.NoError(t, err)
+					queue = append(queue, bucketItem{b: childBucket, bs: child})
+				}
+			}
+		}
+		return nil
+	})
+	require.NoError(t, err)
+
+	t.Log("Inspecting the database")
+	_ = db.View(func(tx *bolt.Tx) error {
+		actualStructure := tx.Inspect()
+		assert.Equal(t, expectedStructure, actualStructure)
+		return nil
+	})
+}
+
 // Ensure a large bucket can calculate stats.
 func TestBucket_Stats_Large(t *testing.T) {
 	if testing.Short() {

cmd/bbolt/README.md

@@ -162,6 +162,61 @@
       Bytes used for inlined buckets: 780 (0%)
   ```
 
+### inspect
+- `inspect` inspect the structure of the database.
+- Usage: `bbolt inspect [path to the bbolt database]`
+
+  Example:
+```bash
+$ ./bbolt inspect ~/default.etcd/member/snap/db
+{
+    "name": "root",
+    "keyN": 0,
+    "buckets": [
+        {
+            "name": "alarm",
+            "keyN": 0
+        },
+        {
+            "name": "auth",
+            "keyN": 2
+        },
+        {
+            "name": "authRoles",
+            "keyN": 1
+        },
+        {
+            "name": "authUsers",
+            "keyN": 1
+        },
+        {
+            "name": "cluster",
+            "keyN": 1
+        },
+        {
+            "name": "key",
+            "keyN": 1285
+        },
+        {
+            "name": "lease",
+            "keyN": 2
+        },
+        {
+            "name": "members",
+            "keyN": 1
+        },
+        {
+            "name": "members_removed",
+            "keyN": 0
+        },
+        {
+            "name": "meta",
+            "keyN": 3
+        }
+    ]
+}
+```
+
 ### pages
 
 - Pages prints a table of pages with their type (meta, leaf, branch, freelist).

cmd/bbolt/command_inspect.go

@@ -0,0 +1,55 @@
+package main
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+
+	bolt "go.etcd.io/bbolt"
+)
+
+func newInspectCobraCommand() *cobra.Command {
+	inspectCmd := &cobra.Command{
+		Use:   "inspect",
+		Short: "inspect the structure of the database",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return errors.New("db file path not provided")
+			}
+			if len(args) > 1 {
+				return errors.New("too many arguments")
+			}
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return inspectFunc(args[0])
+		},
+	}
+
+	return inspectCmd
+}
+
+func inspectFunc(srcDBPath string) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
+	db, err := bolt.Open(srcDBPath, 0600, &bolt.Options{ReadOnly: true})
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	return db.View(func(tx *bolt.Tx) error {
+		bs := tx.Inspect()
+		out, err := json.MarshalIndent(bs, "", "    ")
+		if err != nil {
+			return err
+		}
+		fmt.Fprintln(os.Stdout, string(out))
+		return nil
+	})
+}

cmd/bbolt/command_inspect_test.go

@@ -0,0 +1,27 @@
+package main_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	bolt "go.etcd.io/bbolt"
+	main "go.etcd.io/bbolt/cmd/bbolt"
+	"go.etcd.io/bbolt/internal/btesting"
+)
+
+func TestInspect(t *testing.T) {
+	pageSize := 4096
+	db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
+	srcPath := db.Path()
+	db.Close()
+
+	defer requireDBNoChange(t, dbData(t, db.Path()), db.Path())
+
+	rootCmd := main.NewRootCommand()
+	rootCmd.SetArgs([]string{
+		"inspect", srcPath,
+	})
+	err := rootCmd.Execute()
+	require.NoError(t, err)
+}

cmd/bbolt/command_root.go

@@ -19,6 +19,7 @@ func NewRootCommand() *cobra.Command {
 	rootCmd.AddCommand(
 		newVersionCobraCommand(),
 		newSurgeryCobraCommand(),
+		newInspectCobraCommand(),
 	)
 
 	return rootCmd

cmd/bbolt/command_surgery.go

@@ -330,13 +330,3 @@ func readMetaPage(path string) (*common.Meta, error) {
 	}
 	return m[1], nil
 }
-
-func checkSourceDBPath(srcPath string) (os.FileInfo, error) {
-	fi, err := os.Stat(srcPath)
-	if os.IsNotExist(err) {
-		return nil, fmt.Errorf("source database file %q doesn't exist", srcPath)
-	} else if err != nil {
-		return nil, fmt.Errorf("failed to open source database file %q: %v", srcPath, err)
-	}
-	return fi, nil
-}

cmd/bbolt/main.go

@@ -170,6 +170,7 @@ The commands are:
     pages       print list of pages with their types
     page-item   print the key and value of a page item.
     stats       iterate over all pages and generate usage stats
+    inspect     inspect the structure of the database
     surgery     perform surgery on bbolt database
 
 Use "bbolt [command] -h" for more information about a command.

cmd/bbolt/utils.go

@@ -0,0 +1,16 @@
+package main
+
+import (
+	"fmt"
+	"os"
+)
+
+func checkSourceDBPath(srcPath string) (os.FileInfo, error) {
+	fi, err := os.Stat(srcPath)
+	if os.IsNotExist(err) {
+		return nil, fmt.Errorf("source database file %q doesn't exist", srcPath)
+	} else if err != nil {
+		return nil, fmt.Errorf("failed to open source database file %q: %v", srcPath, err)
+	}
+	return fi, nil
+}

tx.go

@@ -100,6 +100,11 @@ func (tx *Tx) Stats() TxStats {
 	return tx.stats
 }
 
+// Inspect returns the structure of the database.
+func (tx *Tx) Inspect() BucketStructure {
+	return tx.root.Inspect()
+}
+
 // Bucket retrieves a bucket by name.
 // Returns nil if the bucket does not exist.
 // The bucket instance is only valid for the lifetime of the transaction.