From 10fed5f74d44caad8f777a38db735ebb7fb3e65b Mon Sep 17 00:00:00 2001 From: Ben Johnson Date: Fri, 11 Apr 2014 14:59:46 -0600 Subject: [PATCH] Upgrade import/export to use nested buckets. --- cmd/bolt/buckets_test.go | 6 +++--- cmd/bolt/export.go | 28 +++++++++++++++++++++++----- cmd/bolt/export_test.go | 17 ++++++++++++----- cmd/bolt/get_test.go | 6 +++--- cmd/bolt/import.go | 27 ++++++++++++++++++++++++--- cmd/bolt/import_test.go | 11 ++++++++--- cmd/bolt/keys_test.go | 8 ++++---- cmd/bolt/set_test.go | 2 +- 8 files changed, 78 insertions(+), 27 deletions(-) diff --git a/cmd/bolt/buckets_test.go b/cmd/bolt/buckets_test.go index 5f72bb2..27ee619 100644 --- a/cmd/bolt/buckets_test.go +++ b/cmd/bolt/buckets_test.go @@ -13,9 +13,9 @@ func TestBuckets(t *testing.T) { SetTestMode(true) open(func(db *bolt.DB, path string) { db.Update(func(tx *bolt.Tx) error { - tx.CreateBucket("woojits") - tx.CreateBucket("widgets") - tx.CreateBucket("whatchits") + tx.CreateBucket([]byte("woojits")) + tx.CreateBucket([]byte("widgets")) + tx.CreateBucket([]byte("whatchits")) return nil }) db.Close() diff --git a/cmd/bolt/export.go b/cmd/bolt/export.go index f3cafc1..2dcbc1f 100644 --- a/cmd/bolt/export.go +++ b/cmd/bolt/export.go @@ -23,25 +23,33 @@ func Export(path string) { } defer db.Close() - db.View(func(tx *bolt.Tx) error { + err = db.View(func(tx *bolt.Tx) error { // Loop over every bucket and export it as a raw message. var root []*rawMessage - for _, b := range tx.Buckets() { + err := tx.ForEach(func(name []byte, b *bolt.Bucket) error { message, err := exportBucket(b) if err != nil { fatal(err) } + message.Key = name root = append(root, message) + return nil + }) + if err != nil { + return err } // Encode all buckets into JSON. output, err := json.Marshal(root) if err != nil { - fatal("encode: ", err) + return fmt.Errorf("encode: ", err) } print(string(output)) return nil }) + if err != nil { + fatal(err) + } } func exportBucket(b *bolt.Bucket) (*rawMessage, error) { @@ -50,11 +58,22 @@ func exportBucket(b *bolt.Bucket) (*rawMessage, error) { err := b.ForEach(func(k, v []byte) error { var err error + // If there is no value then it is a bucket. + if v == nil { + child, err := exportBucket(b.Bucket(k)) + if err != nil { + return fmt.Errorf("bucket: %s", err) + } + child.Key = k + children = append(children, child) + return nil + } + + // Otherwise it's a regular key. var child = &rawMessage{Key: k} if child.Value, err = json.Marshal(v); err != nil { return fmt.Errorf("value: %s", err) } - children = append(children, child) return nil }) @@ -64,7 +83,6 @@ func exportBucket(b *bolt.Bucket) (*rawMessage, error) { // Encode bucket into a raw message. var root = rawMessage{Type: "bucket"} - root.Key = []byte(b.Name()) if root.Value, err = json.Marshal(children); err != nil { return nil, fmt.Errorf("children: %s", err) } diff --git a/cmd/bolt/export_test.go b/cmd/bolt/export_test.go index 3d6c21a..13f57d1 100644 --- a/cmd/bolt/export_test.go +++ b/cmd/bolt/export_test.go @@ -13,19 +13,26 @@ func TestExport(t *testing.T) { SetTestMode(true) open(func(db *bolt.DB, path string) { db.Update(func(tx *bolt.Tx) error { - tx.CreateBucket("widgets") - b := tx.Bucket("widgets") + tx.CreateBucket([]byte("widgets")) + b := tx.Bucket([]byte("widgets")) b.Put([]byte("foo"), []byte("0000")) b.Put([]byte("bar"), []byte("")) - tx.CreateBucket("woojits") - b = tx.Bucket("woojits") + tx.CreateBucket([]byte("woojits")) + b = tx.Bucket([]byte("woojits")) b.Put([]byte("baz"), []byte("XXXX")) + + b.CreateBucket([]byte("woojits/subbucket")) + b = b.Bucket([]byte("woojits/subbucket")) + b.Put([]byte("bat"), []byte("A")) + + tx.CreateBucket([]byte("empty")) + return nil }) db.Close() output := run("export", path) - assert.Equal(t, `[{"type":"bucket","key":"d2lkZ2V0cw==","value":[{"key":"YmFy","value":""},{"key":"Zm9v","value":"MDAwMA=="}]},{"type":"bucket","key":"d29vaml0cw==","value":[{"key":"YmF6","value":"WFhYWA=="}]}]`, output) + assert.Equal(t, `[{"type":"bucket","key":"ZW1wdHk=","value":[]},{"type":"bucket","key":"d2lkZ2V0cw==","value":[{"key":"YmFy","value":""},{"key":"Zm9v","value":"MDAwMA=="}]},{"type":"bucket","key":"d29vaml0cw==","value":[{"key":"YmF6","value":"WFhYWA=="},{"type":"bucket","key":"d29vaml0cy9zdWJidWNrZXQ=","value":[{"key":"YmF0","value":"QQ=="}]}]}]`, output) }) } diff --git a/cmd/bolt/get_test.go b/cmd/bolt/get_test.go index 09883d4..d491971 100644 --- a/cmd/bolt/get_test.go +++ b/cmd/bolt/get_test.go @@ -13,8 +13,8 @@ func TestGet(t *testing.T) { SetTestMode(true) open(func(db *bolt.DB, path string) { db.Update(func(tx *bolt.Tx) error { - tx.CreateBucket("widgets") - tx.Bucket("widgets").Put([]byte("foo"), []byte("bar")) + tx.CreateBucket([]byte("widgets")) + tx.Bucket([]byte("widgets")).Put([]byte("foo"), []byte("bar")) return nil }) db.Close() @@ -45,7 +45,7 @@ func TestGetKeyNotFound(t *testing.T) { SetTestMode(true) open(func(db *bolt.DB, path string) { db.Update(func(tx *bolt.Tx) error { - return tx.CreateBucket("widgets") + return tx.CreateBucket([]byte("widgets")) }) db.Close() output := run("get", path, "widgets", "foo") diff --git a/cmd/bolt/import.go b/cmd/bolt/import.go index ec8cee1..7554ae7 100644 --- a/cmd/bolt/import.go +++ b/cmd/bolt/import.go @@ -41,7 +41,7 @@ func Import(path string, input string) { } // Create the bucket if it doesn't exist. - if err := tx.CreateBucketIfNotExists(string(message.Key)); err != nil { + if err := tx.CreateBucketIfNotExists(message.Key); err != nil { return fmt.Errorf("create bucket: %s", err) } @@ -52,7 +52,7 @@ func Import(path string, input string) { } // Import all the values into the bucket. - b := tx.Bucket(string(message.Key)) + b := tx.Bucket(message.Key) if err := importBucket(b, children); err != nil { return fmt.Errorf("import bucket: %s", err) } @@ -67,7 +67,28 @@ func Import(path string, input string) { func importBucket(b *bolt.Bucket, children []*rawMessage) error { // Decode each message into a key/value pair. for _, child := range children { - // Decode the base64 value. + // Bucket messages are handled recursively. + if child.Type == "bucket" { + // Create the bucket if it doesn't exist. + if err := b.CreateBucketIfNotExists(child.Key); err != nil { + return fmt.Errorf("create bucket: %s", err) + } + + // Decode child messages. + var subchildren []*rawMessage + if err := json.Unmarshal(child.Value, &subchildren); err != nil { + return fmt.Errorf("decode children: %s", err) + } + + // Import subbucket. + subbucket := b.Bucket(child.Key) + if err := importBucket(subbucket, subchildren); err != nil { + return fmt.Errorf("import bucket: %s", err) + } + continue + } + + // Non-bucket values are decoded from base64. var value []byte if err := json.Unmarshal(child.Value, &value); err != nil { return fmt.Errorf("decode value: %s", err) diff --git a/cmd/bolt/import_test.go b/cmd/bolt/import_test.go index be41f5c..263f561 100644 --- a/cmd/bolt/import_test.go +++ b/cmd/bolt/import_test.go @@ -15,7 +15,7 @@ func TestImport(t *testing.T) { // Write input file. input := tempfile() - assert.NoError(t, ioutil.WriteFile(input, []byte(`[{"type":"bucket","key":"d2lkZ2V0cw==","value":[{"key":"YmFy","value":""},{"key":"Zm9v","value":"MDAwMA=="}]},{"type":"bucket","key":"d29vaml0cw==","value":[{"key":"YmF6","value":"WFhYWA=="}]}]`), 0600)) + assert.NoError(t, ioutil.WriteFile(input, []byte(`[{"type":"bucket","key":"ZW1wdHk=","value":[]},{"type":"bucket","key":"d2lkZ2V0cw==","value":[{"key":"YmFy","value":""},{"key":"Zm9v","value":"MDAwMA=="}]},{"type":"bucket","key":"d29vaml0cw==","value":[{"key":"YmF6","value":"WFhYWA=="},{"type":"bucket","key":"d29vaml0cy9zdWJidWNrZXQ=","value":[{"key":"YmF0","value":"QQ=="}]}]}]`), 0600)) // Import database. path := tempfile() @@ -26,15 +26,20 @@ func TestImport(t *testing.T) { db, err := bolt.Open(path, 0600) assert.NoError(t, err) db.View(func(tx *bolt.Tx) error { - b := tx.Bucket("widgets") + assert.NotNil(t, tx.Bucket([]byte("empty"))) + + b := tx.Bucket([]byte("widgets")) if assert.NotNil(t, b) { assert.Equal(t, []byte("0000"), b.Get([]byte("foo"))) assert.Equal(t, []byte(""), b.Get([]byte("bar"))) } - b = tx.Bucket("woojits") + b = tx.Bucket([]byte("woojits")) if assert.NotNil(t, b) { assert.Equal(t, []byte("XXXX"), b.Get([]byte("baz"))) + + b = b.Bucket([]byte("woojits/subbucket")) + assert.Equal(t, []byte("A"), b.Get([]byte("bat"))) } return nil diff --git a/cmd/bolt/keys_test.go b/cmd/bolt/keys_test.go index ea530f6..2b5a9a0 100644 --- a/cmd/bolt/keys_test.go +++ b/cmd/bolt/keys_test.go @@ -13,10 +13,10 @@ func TestKeys(t *testing.T) { SetTestMode(true) open(func(db *bolt.DB, path string) { db.Update(func(tx *bolt.Tx) error { - tx.CreateBucket("widgets") - tx.Bucket("widgets").Put([]byte("0002"), []byte("")) - tx.Bucket("widgets").Put([]byte("0001"), []byte("")) - tx.Bucket("widgets").Put([]byte("0003"), []byte("")) + tx.CreateBucket([]byte("widgets")) + tx.Bucket([]byte("widgets")).Put([]byte("0002"), []byte("")) + tx.Bucket([]byte("widgets")).Put([]byte("0001"), []byte("")) + tx.Bucket([]byte("widgets")).Put([]byte("0003"), []byte("")) return nil }) db.Close() diff --git a/cmd/bolt/set_test.go b/cmd/bolt/set_test.go index 519d888..be07148 100644 --- a/cmd/bolt/set_test.go +++ b/cmd/bolt/set_test.go @@ -13,7 +13,7 @@ func TestSet(t *testing.T) { SetTestMode(true) open(func(db *bolt.DB, path string) { db.Update(func(tx *bolt.Tx) error { - tx.CreateBucket("widgets") + tx.CreateBucket([]byte("widgets")) return nil }) db.Close()