Merge pull request #2 from piax93/bucket-selection-option

add option to select only certain buckets

Author: konoui

README.md

@@ -8,7 +8,7 @@ $ go get -u github.com/konoui/boltdb-exporter
 
 ## Usage
 ```
-$ boltdb-exporter --db <database filename> --format yaml
+$ boltdb-exporter --db <database filename> --format yaml [--bucket <root bucket name> ...]
 ```
 
 ## Example

main.go

@@ -12,10 +12,22 @@ import (
 	"github.com/peterbourgon/ff/v3/ffcli"
 )
 
+type arrayFlag []string
+
 type config struct {
-	filename     string
-	outputFormat string
-	marshaler    func(interface{}) ([]byte, error)
+	filename        string
+	outputFormat    string
+	bucketSelection arrayFlag
+	marshaler       func(interface{}) ([]byte, error)
+}
+
+func (i *arrayFlag) Set(value string) error {
+	*i = append(*i, value)
+	return nil
+}
+
+func (i *arrayFlag) String() string {
+	return "..."
 }
 
 func newRootCmd() *ffcli.Command {
@@ -37,6 +49,7 @@ func newRootCmd() *ffcli.Command {
 func (cfg *config) registerFlags(fs *flag.FlagSet) {
 	fs.StringVar(&cfg.outputFormat, "format", "json", "support json/yaml")
 	fs.StringVar(&cfg.filename, "db", "", "database filename")
+	fs.Var(&cfg.bucketSelection, "bucket", "select root-level bucket to export (can be used multiple times)")
 }
 
 func (cfg *config) validate() error {
@@ -66,7 +79,15 @@ func (cfg *config) run() error {
 		return err
 	}
 
-	b, err := exporter.Export(cfg.filename, cfg.marshaler)
+	var bucketSelectionSet map[string]bool
+	if len(cfg.bucketSelection) > 0 {
+		bucketSelectionSet = make(map[string]bool)
+		for _, bucket := range cfg.bucketSelection {
+			bucketSelectionSet[bucket] = true
+		}
+	}
+
+	b, err := exporter.Export(cfg.filename, cfg.marshaler, bucketSelectionSet)
 	if err != nil {
 		return err
 	}

pkg/exporter/export.go

@@ -6,21 +6,21 @@ import (
 	"github.com/boltdb/bolt"
 )
 
-func Export(filename string, marshaler func(interface{}) ([]byte, error)) (ret []byte, err error) {
+func Export(filename string, marshaler func(interface{}) ([]byte, error), bucketSelection map[string]bool) (ret []byte, err error) {
 	db, err := bolt.Open(filename, 0600, &bolt.Options{
 		ReadOnly: true,
 	})
 	if err != nil {
 		return nil, err
 	}
 	defer db.Close()
-	return export(db, marshaler)
+	return export(db, marshaler, bucketSelection)
 }
 
-func export(db *bolt.DB, marshaler func(interface{}) ([]byte, error)) (ret []byte, err error) {
+func export(db *bolt.DB, marshaler func(interface{}) ([]byte, error), bucketSelection map[string]bool) (ret []byte, err error) {
 	err = db.View(func(tx *bolt.Tx) error {
 		c := tx.Cursor()
-		rawMap := makeRawMap(tx, c)
+		rawMap := makeRawMap(tx, c, bucketSelection)
 		ret, err = marshaler(rawMap)
 		if err != nil {
 			return err
@@ -30,22 +30,29 @@ func export(db *bolt.DB, marshaler func(interface{}) ([]byte, error)) (ret []byt
 	return
 }
 
-func makeRawMap(tx *bolt.Tx, c *bolt.Cursor) map[string]interface{} {
+func makeRawMap(tx *bolt.Tx, c *bolt.Cursor, bucketSelection map[string]bool) map[string]interface{} {
 	rawMap := make(map[string]interface{})
-	recursiveRawMap(tx, c, rawMap)
+	recursiveRawMap(tx, c, rawMap, bucketSelection)
 	return rawMap
 }
 
-func recursiveRawMap(tx *bolt.Tx, c *bolt.Cursor, rawMap map[string]interface{}) map[string]interface{} {
+func recursiveRawMap(tx *bolt.Tx, c *bolt.Cursor, rawMap map[string]interface{}, bucketSelection map[string]bool) map[string]interface{} {
 	for k, v := c.First(); k != nil; k, v = c.Next() {
+		// skip bucket if not selected
+		if bucketSelection != nil {
+			if _, ok := bucketSelection[string(k)]; !ok {
+				continue
+			}
+		}
+
 		if v == nil {
 			bucket := c.Bucket().Bucket(k)
 			if bucket == nil {
 				bucket = tx.Bucket(k)
 			}
 			nextCursor := bucket.Cursor()
 			nextMap := make(map[string]interface{})
-			rawMap[string(k)] = recursiveRawMap(tx, nextCursor, nextMap)
+			rawMap[string(k)] = recursiveRawMap(tx, nextCursor, nextMap, nil)
 			continue
 		}
 

pkg/exporter/export_test.go

@@ -45,20 +45,29 @@ var testData = TestData{
 
 func Test_export(t *testing.T) {
 	tests := []struct {
-		name     string
-		wantPath string
-		setuper  func(*testing.T, string) *bolt.DB
-		update   bool
+		name      string
+		wantPath  string
+		setuper   func(*testing.T, string) *bolt.DB
+		update    bool
+		selection map[string]bool
 	}{
 		{
-			name:     "multi buckets",
-			wantPath: filepath.Join("testdata", "single-output.json"),
-			setuper:  setupSingle,
+			name:      "multi buckets",
+			wantPath:  filepath.Join("testdata", "single-output.json"),
+			setuper:   setupSingle,
+			selection: nil,
 		},
 		{
-			name:     "nested buckets",
-			setuper:  setupNested,
-			wantPath: filepath.Join("testdata", "nested-output.json"),
+			name:      "nested buckets",
+			setuper:   setupNested,
+			wantPath:  filepath.Join("testdata", "nested-output.json"),
+			selection: nil,
+		},
+		{
+			name:      "bucket selection",
+			wantPath:  filepath.Join("testdata", "single-output-selection.json"),
+			setuper:   setupSingle,
+			selection: map[string]bool{"bucket2": true},
 		},
 	}
 
@@ -67,7 +76,7 @@ func Test_export(t *testing.T) {
 			db := tt.setuper(t, filename)
 			defer cleanup(t, db, filename)
 
-			gotData, err := export(db, json.Marshal)
+			gotData, err := export(db, json.Marshal, tt.selection)
 			if err != nil {
 				t.Error(err)
 			}

pkg/exporter/testdata/single-output-selection.json

@@ -0,0 +1,54 @@
+{
+    "bucket2": {
+        "1": {
+            "Data1": "1",
+            "Data2": "2",
+            "Data3": 3,
+            "Data4": {
+                "Data1": "4-1",
+                "Data2": "4-2",
+                "Data3": 3,
+                "Data4": null,
+                "Data5": null
+            },
+            "Data5": {
+                "Data1": "data5-nested1",
+                "Data2": "data5-nested2",
+                "Data3": 0,
+                "Data4": {
+                    "Data1": "5-4-1",
+                    "Data2": "5-4-2",
+                    "Data3": 541,
+                    "Data4": null,
+                    "Data5": null
+                },
+                "Data5": null
+            }
+        },
+        "2": {
+            "Data1": "1",
+            "Data2": "2",
+            "Data3": 3,
+            "Data4": {
+                "Data1": "4-1",
+                "Data2": "4-2",
+                "Data3": 3,
+                "Data4": null,
+                "Data5": null
+            },
+            "Data5": {
+                "Data1": "data5-nested1",
+                "Data2": "data5-nested2",
+                "Data3": 0,
+                "Data4": {
+                    "Data1": "5-4-1",
+                    "Data2": "5-4-2",
+                    "Data3": 541,
+                    "Data4": null,
+                    "Data5": null
+                },
+                "Data5": null
+            }
+        }
+    }
+}