6 files changed, 67 insertions, 50 deletions
diff --git a/cmd/bolt/bench.go b/cmd/bolt/bench.go
index 28cbbc3..f894a70 100644
--- a/cmd/bolt/bench.go
+++ b/cmd/bolt/bench.go
@@ -35,7 +35,15 @@ func Bench(inputPath string, readWriteMode string, traversalPattern string, para
 
 	// benchmarks for getting all keys
 
-	b := bolt.NewBenchmark(inputPath, readWriteMode, traversalPattern, parallelism)
+	// Open the database.
+	db, err := bolt.Open(inputPath, 0600)
+	if err != nil {
+		fatalf("error: %+v", err)
+		return
+	}
+	defer db.Close()
+
+	b := bolt.NewBenchmark(db, readWriteMode, traversalPattern, parallelism)
 
 	result := testing.Benchmark(b.Run)
 
diff --git a/cmd/bolt/generate.go b/cmd/bolt/generate.go
index ecd391e..15edb27 100644
--- a/cmd/bolt/generate.go
+++ b/cmd/bolt/generate.go
@@ -1,32 +1,55 @@
 package main
 
 import (
-	"bufio"
 	"fmt"
-	"os"
-	"strings"
+
+	"github.com/boltdb/bolt"
 )
 
 // Generate data for benchmarks.
-func Generate(numEvents int, destPath string) {
-	f, err := os.Create(destPath)
+func Generate(destPath string, numBuckets, numItems int) {
+
+	// Open the database.
+	db, err := bolt.Open(destPath, 0600)
 	if err != nil {
-		fatal(err)
+		fatalf("open db:", err)
+		return
 	}
-	defer func() {
-		if err := f.Close(); err != nil {
-			fatal(err)
-		}
-	}()
-	w := bufio.NewWriter(f)
+	defer db.Close()
 
-	for i := 0; i < numEvents; i++ {
-		if _, err := w.Write([]byte(fmt.Sprintf("key%d:%s\n", i, strings.Repeat("0", 64)))); err != nil {
-			fatal(err)
-		}
-	}
+	for bucketIndex := 0; bucketIndex < numBuckets; bucketIndex++ {
+		bucketName := fmt.Sprintf("bucket%03d", bucketIndex)
+
+		err = db.Update(func(tx *bolt.Tx) error {
+
+			// Create the bucket if it doesn't exist.
+			if err := tx.CreateBucketIfNotExists([]byte(bucketName)); err != nil {
+				fatalf("create bucket: %s", err)
+				return nil
+			}
 
-	if err = w.Flush(); err != nil {
+			// Find bucket.
+			b := tx.Bucket([]byte(bucketName))
+			if b == nil {
+				fatalf("bucket not found: %s", bucketName)
+				return nil
+			}
+
+			for i := 0; i < numItems; i++ {
+				key := fmt.Sprintf("key%03d", i)
+				value := fmt.Sprintf("value%03d", i)
+
+				// Set value for a given key.
+				if err := b.Put([]byte(key), []byte(value)); err != nil {
+					return err
+				}
+			}
+
+			return nil
+		})
+	}
+	if err != nil {
 		fatal(err)
+		return
 	}
 }
diff --git a/cmd/bolt/generate_test.go b/cmd/bolt/generate_test.go
deleted file mode 100644
index 06ab7d0..0000000
--- a/cmd/bolt/generate_test.go
+++ /dev/null
@@ -1 +0,0 @@
-package main
diff --git a/cmd/bolt/keys.go b/cmd/bolt/keys.go
index 65b717f..6affefe 100644
--- a/cmd/bolt/keys.go
+++ b/cmd/bolt/keys.go
@@ -1,8 +1,6 @@
 package main
 
 import (
-	"errors"
-	"fmt"
 	"os"
 
 	"github.com/boltdb/bolt"
@@ -10,44 +8,34 @@ import (
 
 // Keys retrieves a list of keys for a given bucket.
 func Keys(path, name string) {
-	keys, err := keys(path, name)
-
-	if err != nil {
+	if _, err := os.Stat(path); os.IsNotExist(err) {
 		fatal(err)
 		return
 	}
 
-	for _, key := range keys {
-		println(key)
-	}
-}
-
-func keys(path, name string) ([]string, error) {
-	if _, err := os.Stat(path); os.IsNotExist(err) {
-		return nil, err
-	}
-
 	db, err := bolt.Open(path, 0600)
 	if err != nil {
-		return nil, err
+		fatal(err)
+		return
 	}
 	defer db.Close()
 
-	keys := []string{}
-
 	err = db.View(func(tx *bolt.Tx) error {
 		// Find bucket.
 		b := tx.Bucket([]byte(name))
 		if b == nil {
-			return errors.New(fmt.Sprintf("bucket %+v not found", b))
+			fatalf("bucket not found: %s", name)
+			return nil
 		}
 
 		// Iterate over each key.
 		return b.ForEach(func(key, _ []byte) error {
-			keys = append(keys, string(key))
+			println(string(key))
 			return nil
 		})
 	})
-
-	return keys, err
+	if err != nil {
+		fatal(err)
+		return
+	}
 }
diff --git a/cmd/bolt/main.go b/cmd/bolt/main.go
index 19df59f..ac71631 100644
--- a/cmd/bolt/main.go
+++ b/cmd/bolt/main.go
@@ -95,12 +95,16 @@ func NewApp() *cli.App {
 			Name:  "generate",
 			Usage: "Generate data for benchmarks",
 			Action: func(c *cli.Context) {
-				numEvents, err := strconv.Atoi(c.Args().Get(0))
+				destPath := c.Args().Get(0)
+				numBuckets, err := strconv.Atoi(c.Args().Get(1))
 				if err != nil {
 					fatal(err)
 				}
-				destPath := c.Args().Get(1)
-				Generate(numEvents, destPath)
+				numItems, err := strconv.Atoi(c.Args().Get(2))
+				if err != nil {
+					fatal(err)
+				}
+				Generate(destPath, numBuckets, numItems)
 			},
 		},
 		{
diff --git a/cmd/bolt/set.go b/cmd/bolt/set.go
index c757d27..f4a4696 100644
--- a/cmd/bolt/set.go
+++ b/cmd/bolt/set.go
@@ -21,11 +21,6 @@ func Set(path, name, key, value string) {
 	defer db.Close()
 
 	err = db.Update(func(tx *bolt.Tx) error {
-		// Create the bucket if it doesn't exist.
-		if err := tx.CreateBucketIfNotExists([]byte(name)); err != nil {
-			fatalf("create bucket: %s", err)
-			return nil
-		}
 
 		// Find bucket.
 		b := tx.Bucket([]byte(name))