Merge pull request #128 from benbjohnson/import-export

Add import/export to CLI.

Makefile

@@ -42,6 +42,10 @@ test: fmt errcheck
 	@go test -v -cover -test.run=$(TEST)
 	@echo ""
 	@echo ""
+	@echo "=== CLI ==="
+	@go test -v -test.run=$(TEST) ./cmd/bolt
+	@echo ""
+	@echo ""
 	@echo "=== RACE DETECTOR ==="
 	@go test -v -race -test.run=Parallel
 

cmd/bolt/buckets_test.go

@@ -11,14 +11,15 @@ import (
 // Ensure that a list of buckets can be retrieved.
 func TestBuckets(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	open(func(db *bolt.DB, path string) {
 		db.Update(func(tx *bolt.Tx) error {
 			tx.CreateBucket("woojits")
 			tx.CreateBucket("widgets")
 			tx.CreateBucket("whatchits")
 			return nil
 		})
-		output := run("buckets", db.Path())
+		db.Close()
+		output := run("buckets", path)
 		assert.Equal(t, "whatchits\nwidgets\nwoojits", output)
 	})
 }

cmd/bolt/export.go

@@ -0,0 +1,73 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+
+	"github.com/boltdb/bolt"
+)
+
+// Export exports the entire database as a JSON document.
+func Export(path string) {
+	if _, err := os.Stat(path); os.IsNotExist(err) {
+		fatal(err)
+		return
+	}
+
+	// Open the database.
+	db, err := bolt.Open(path, 0600)
+	if err != nil {
+		fatal(err)
+		return
+	}
+	defer db.Close()
+
+	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() {
+			message, err := exportBucket(b)
+			if err != nil {
+				fatal(err)
+			}
+			root = append(root, message)
+		}
+
+		// Encode all buckets into JSON.
+		output, err := json.Marshal(root)
+		if err != nil {
+			fatal("encode: ", err)
+		}
+		print(string(output))
+		return nil
+	})
+}
+
+func exportBucket(b *bolt.Bucket) (*rawMessage, error) {
+	// Encode individual key/value pairs into raw messages.
+	var children = make([]*rawMessage, 0)
+	err := b.ForEach(func(k, v []byte) error {
+		var err error
+
+		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
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	// 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)
+	}
+
+	return &root, nil
+}

cmd/bolt/export_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"
+)
+
+// Ensure that a database can be exported.
+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")
+			b.Put([]byte("foo"), []byte("0000"))
+			b.Put([]byte("bar"), []byte(""))
+
+			tx.CreateBucket("woojits")
+			b = tx.Bucket("woojits")
+			b.Put([]byte("baz"), []byte("XXXX"))
+			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)
+	})
+}
+
+// Ensure that an error is reported if the database is not found.
+func TestExport_NotFound(t *testing.T) {
+	SetTestMode(true)
+	output := run("export", "no/such/db")
+	assert.Equal(t, "stat no/such/db: no such file or directory", output)
+}

cmd/bolt/get_test.go

@@ -11,13 +11,14 @@ import (
 // Ensure that a value can be retrieved from the CLI.
 func TestGet(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	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"))
 			return nil
 		})
-		output := run("get", db.Path(), "widgets", "foo")
+		db.Close()
+		output := run("get", path, "widgets", "foo")
 		assert.Equal(t, "bar", output)
 	})
 }
@@ -32,8 +33,9 @@ func TestGetDBNotFound(t *testing.T) {
 // Ensure that an error is reported if the bucket is not found.
 func TestGetBucketNotFound(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	open(func(db *bolt.DB, path string) {
-		output := run("get", db.Path(), "widgets", "foo")
+		db.Close()
+		output := run("get", path, "widgets", "foo")
 		assert.Equal(t, "bucket not found: widgets", output)
 	})
 }
@@ -41,11 +43,12 @@ func TestGetBucketNotFound(t *testing.T) {
 // Ensure that an error is reported if the key is not found.
 func TestGetKeyNotFound(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	open(func(db *bolt.DB, path string) {
 		db.Update(func(tx *bolt.Tx) error {
 			return tx.CreateBucket("widgets")
 		})
-		output := run("get", db.Path(), "widgets", "foo")
+		db.Close()
+		output := run("get", path, "widgets", "foo")
 		assert.Equal(t, "key not found: foo", output)
 	})
 }

cmd/bolt/import.go

@@ -0,0 +1,82 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+
+	"github.com/boltdb/bolt"
+)
+
+// Import converts an exported database dump into a new database.
+func Import(path string, input string) {
+	f, err := os.Open(input)
+	if err != nil {
+		fatal(err)
+		return
+	}
+	defer f.Close()
+
+	// Read in entire dump.
+	var root []*rawMessage
+	if err := json.NewDecoder(f).Decode(&root); err != nil {
+		fatal(err)
+	}
+
+	// Open the database.
+	db, err := bolt.Open(path, 0600)
+	if err != nil {
+		fatal(err)
+		return
+	}
+	defer db.Close()
+
+	// Insert entire dump into database.
+	err = db.Update(func(tx *bolt.Tx) error {
+		// Loop over every message and create a bucket.
+		for _, message := range root {
+			// Validate that root messages are buckets.
+			if message.Type != "bucket" {
+				return fmt.Errorf("invalid root type: %q", message.Type)
+			}
+
+			// Create the bucket if it doesn't exist.
+			if err := tx.CreateBucketIfNotExists(string(message.Key)); err != nil {
+				return fmt.Errorf("create bucket: %s", err)
+			}
+
+			// Decode child messages.
+			var children []*rawMessage
+			if err := json.Unmarshal(message.Value, &children); err != nil {
+				return fmt.Errorf("decode children: %s", err)
+			}
+
+			// Import all the values into the bucket.
+			b := tx.Bucket(string(message.Key))
+			if err := importBucket(b, children); err != nil {
+				return fmt.Errorf("import bucket: %s", err)
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		fatal("update: ", err)
+	}
+}
+
+func importBucket(b *bolt.Bucket, children []*rawMessage) error {
+	// Decode each message into a key/value pair.
+	for _, child := range children {
+		// Decode the base64 value.
+		var value []byte
+		if err := json.Unmarshal(child.Value, &value); err != nil {
+			return fmt.Errorf("decode value: %s", err)
+		}
+
+		// Insert key/value into bucket.
+		if err := b.Put(child.Key, value); err != nil {
+			return fmt.Errorf("put: %s", err)
+		}
+	}
+	return nil
+}

cmd/bolt/import_test.go

@@ -0,0 +1,50 @@
+package main_test
+
+import (
+	"io/ioutil"
+	"testing"
+
+	"github.com/boltdb/bolt"
+	. "github.com/boltdb/bolt/cmd/bolt"
+	"github.com/stretchr/testify/assert"
+)
+
+// Ensure that a database can be imported.
+func TestImport(t *testing.T) {
+	SetTestMode(true)
+
+	// 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))
+
+	// Import database.
+	path := tempfile()
+	output := run("import", path, "--input", input)
+	assert.Equal(t, ``, output)
+
+	// Open database and verify contents.
+	db, err := bolt.Open(path, 0600)
+	assert.NoError(t, err)
+	db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket("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")
+		if assert.NotNil(t, b) {
+			assert.Equal(t, []byte("XXXX"), b.Get([]byte("baz")))
+		}
+
+		return nil
+	})
+	db.Close()
+}
+
+// Ensure that an error is reported if the database is not found.
+func TestImport_NotFound(t *testing.T) {
+	SetTestMode(true)
+	output := run("import", "path/to/db", "--input", "no/such/file")
+	assert.Equal(t, "open no/such/file: no such file or directory", output)
+}

cmd/bolt/keys_test.go

@@ -11,7 +11,7 @@ import (
 // Ensure that a list of keys can be retrieved for a given bucket.
 func TestKeys(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	open(func(db *bolt.DB, path string) {
 		db.Update(func(tx *bolt.Tx) error {
 			tx.CreateBucket("widgets")
 			tx.Bucket("widgets").Put([]byte("0002"), []byte(""))
@@ -19,7 +19,8 @@ func TestKeys(t *testing.T) {
 			tx.Bucket("widgets").Put([]byte("0003"), []byte(""))
 			return nil
 		})
-		output := run("keys", db.Path(), "widgets")
+		db.Close()
+		output := run("keys", path, "widgets")
 		assert.Equal(t, "0001\n0002\n0003", output)
 	})
 }
@@ -34,8 +35,9 @@ func TestKeysDBNotFound(t *testing.T) {
 // Ensure that an error is reported if the bucket is not found.
 func TestKeysBucketNotFound(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	open(func(db *bolt.DB, path string) {
-		output := run("keys", db.Path(), "widgets")
+		db.Close()
+		output := run("keys", path, "widgets")
 		assert.Equal(t, "bucket not found: widgets", output)
 	})
 }

cmd/bolt/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"log"
 	"os"
@@ -55,6 +56,24 @@ func NewApp() *cli.App {
 				Buckets(path)
 			},
 		},
+		{
+			Name:  "import",
+			Usage: "Imports from a JSON dump into a database",
+			Flags: []cli.Flag{
+				&cli.StringFlag{Name: "input"},
+			},
+			Action: func(c *cli.Context) {
+				Import(c.Args().Get(0), c.String("input"))
+			},
+		},
+		{
+			Name:  "export",
+			Usage: "Exports a database to JSON",
+			Action: func(c *cli.Context) {
+				path := c.Args().Get(0)
+				Export(path)
+			},
+		},
 		{
 			Name:  "pages",
 			Usage: "Dumps page information for a database",
@@ -144,3 +163,10 @@ func SetTestMode(value bool) {
 		logger = log.New(os.Stderr, "", 0)
 	}
 }
+
+// rawMessage represents a JSON element in the import/export document.
+type rawMessage struct {
+	Type  string          `json:"type,omitempty"`
+	Key   []byte          `json:"key"`
+	Value json.RawMessage `json:"value"`
+}

cmd/bolt/main_test.go

@@ -10,17 +10,15 @@ import (
 )
 
 // open creates and opens a Bolt database in the temp directory.
-func open(fn func(*bolt.DB)) {
+func open(fn func(*bolt.DB, string)) {
-	f, _ := ioutil.TempFile("", "bolt-")
+	path := tempfile()
-	f.Close()
+	defer os.RemoveAll(path)
-	os.Remove(f.Name())
-	defer os.RemoveAll(f.Name())
 
-	db, err := bolt.Open(f.Name(), 0600)
+	db, err := bolt.Open(path, 0600)
 	if err != nil {
 		panic("db open error: " + err.Error())
 	}
-	fn(db)
+	fn(db, path)
 }
 
 // run executes a command against the CLI and returns the output.
@@ -29,3 +27,11 @@ func run(args ...string) string {
 	NewApp().Run(args)
 	return strings.TrimSpace(LogBuffer())
 }
+
+// tempfile returns a temporary file path.
+func tempfile() string {
+	f, _ := ioutil.TempFile("", "bolt-")
+	f.Close()
+	os.Remove(f.Name())
+	return f.Name()
+}

cmd/bolt/set_test.go

@@ -11,13 +11,14 @@ import (
 // Ensure that a value can be set from the CLI.
 func TestSet(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	open(func(db *bolt.DB, path string) {
 		db.Update(func(tx *bolt.Tx) error {
 			tx.CreateBucket("widgets")
 			return nil
 		})
-		assert.Equal(t, "", run("set", db.Path(), "widgets", "foo", "bar"))
+		db.Close()
-		assert.Equal(t, "bar", run("get", db.Path(), "widgets", "foo"))
+		assert.Equal(t, "", run("set", path, "widgets", "foo", "bar"))
+		assert.Equal(t, "bar", run("get", path, "widgets", "foo"))
 	})
 }
 
@@ -31,8 +32,9 @@ func TestSetDBNotFound(t *testing.T) {
 // Ensure that an error is reported if the bucket is not found.
 func TestSetBucketNotFound(t *testing.T) {
 	SetTestMode(true)
-	open(func(db *bolt.DB) {
+	open(func(db *bolt.DB, path string) {
-		output := run("set", db.Path(), "widgets", "foo", "bar")
+		db.Close()
+		output := run("set", path, "widgets", "foo", "bar")
 		assert.Equal(t, "bucket not found: widgets", output)
 	})
 }