From b22480fd32daa1c3f6542f4158916651610c3d6e Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Sun, 16 Feb 2014 13:51:35 -0700
Subject: [PATCH] Add Transaction.ForEach().

---
 db.go           | 11 ++++++++
 db_test.go      | 67 +++++++++++++++++++++++++++++++++++++++++++++++++
 example_test.go | 24 ++++++++++++++++++
 transaction.go  | 19 ++++++++++++++
 4 files changed, 121 insertions(+)

diff --git a/db.go b/db.go
index a519f73..0276d7a 100644
--- a/db.go
+++ b/db.go
@@ -359,6 +359,17 @@ func (db *DB) Do(fn func(*RWTransaction) error) error {
 	return t.Commit()
 }
 
+// ForEach executes a function for each key/value pair in a bucket.
+// An error is returned if the bucket cannot be found.
+func (db *DB) ForEach(name string, fn func(k, v []byte) error) error {
+	t, err := db.Transaction()
+	if err != nil {
+		return err
+	}
+	defer t.Close()
+	return t.ForEach(name, fn)
+}
+
 // Bucket retrieves a reference to a bucket.
 // This is typically useful for checking the existence of a bucket.
 func (db *DB) Bucket(name string) (*Bucket, error) {
diff --git a/db_test.go b/db_test.go
index 43c119d..ecf2290 100644
--- a/db_test.go
+++ b/db_test.go
@@ -1,6 +1,7 @@
 package bolt
 
 import (
+	"bytes"
 	"io"
 	"io/ioutil"
 	"os"
@@ -215,6 +216,72 @@ func TestDBTransactionBlockWhileClosed(t *testing.T) {
 	})
 }
 
+// Ensure a database can loop over all key/value pairs in a bucket.
+func TestDBForEach(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		db.CreateBucket("widgets")
+		db.Put("widgets", []byte("foo"), []byte("0000"))
+		db.Put("widgets", []byte("baz"), []byte("0001"))
+		db.Put("widgets", []byte("bar"), []byte("0002"))
+
+		var index int
+		err := db.ForEach("widgets", func(k, v []byte) error {
+			switch index {
+			case 0:
+				assert.Equal(t, k, []byte("bar"))
+				assert.Equal(t, v, []byte("0002"))
+			case 1:
+				assert.Equal(t, k, []byte("baz"))
+				assert.Equal(t, v, []byte("0001"))
+			case 2:
+				assert.Equal(t, k, []byte("foo"))
+				assert.Equal(t, v, []byte("0000"))
+			}
+			index++
+			return nil
+		})
+		assert.NoError(t, err)
+		assert.Equal(t, index, 3)
+	})
+}
+
+// Ensure a database can stop iteration early.
+func TestDBForEachShortCircuit(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		db.CreateBucket("widgets")
+		db.Put("widgets", []byte("bar"), []byte("0000"))
+		db.Put("widgets", []byte("baz"), []byte("0000"))
+		db.Put("widgets", []byte("foo"), []byte("0000"))
+
+		var index int
+		err := db.ForEach("widgets", func(k, v []byte) error {
+			index++
+			if bytes.Equal(k, []byte("baz")) {
+				return &Error{"marker", nil}
+			}
+			return nil
+		})
+		assert.Equal(t, err, &Error{"marker", nil})
+		assert.Equal(t, index, 2)
+	})
+}
+
+// Ensure a database returns an error when trying to attempt a for each on a missing bucket.
+func TestDBForEachBucketNotFound(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		err := db.ForEach("widgets", func(k, v []byte) error { return nil })
+		assert.Equal(t, err, ErrBucketNotFound)
+	})
+}
+
+// Ensure a closed database returns an error when executing a for each.
+func TestDBForEachWhileClosed(t *testing.T) {
+	withDB(func(db *DB, path string) {
+		err := db.ForEach("widgets", func(k, v []byte) error { return nil })
+		assert.Equal(t, err, ErrDatabaseNotOpen)
+	})
+}
+
 // Ensure a closed database returns an error when finding a bucket.
 func TestDBBucketWhileClosed(t *testing.T) {
 	withDB(func(db *DB, path string) {
diff --git a/example_test.go b/example_test.go
index 655e283..542109c 100644
--- a/example_test.go
+++ b/example_test.go
@@ -87,6 +87,30 @@ func ExampleDB_Do() {
 	// The value of 'foo' is: bar
 }
 
+func ExampleDB_ForEach() {
+	// Open the database.
+	var db DB
+	db.Open("/tmp/bolt/db_foreach.db", 0666)
+	defer db.Close()
+
+	// Insert data into a bucket.
+	db.CreateBucket("animals")
+	db.Put("animals", []byte("dog"), []byte("fun"))
+	db.Put("animals", []byte("cat"), []byte("lame"))
+	db.Put("animals", []byte("liger"), []byte("awesome"))
+
+	// Iterate over items in sorted key order.
+	db.ForEach("animals", func(k, v []byte) error {
+		fmt.Printf("A %s is %s.\n", string(k), string(v))
+		return nil
+	})
+
+	// Output:
+	// A cat is lame.
+	// A dog is fun.
+	// A liger is awesome.
+}
+
 func ExampleRWTransaction() {
 	// Open the database.
 	var db DB
diff --git a/transaction.go b/transaction.go
index 0c9da7b..54aae8c 100644
--- a/transaction.go
+++ b/transaction.go
@@ -93,6 +93,25 @@ func (t *Transaction) Get(name string, key []byte) (value []byte, err error) {
 	return c.Get(key), nil
 }
 
+// ForEach executes a function for each key/value pair in a bucket.
+// An error is returned if the bucket cannot be found.
+func (t *Transaction) ForEach(name string, fn func(k, v []byte) error) error {
+	// Open a cursor on the bucket.
+	c, err := t.Cursor(name)
+	if err != nil {
+		return err
+	}
+
+	// Iterate over each key/value pair in the bucket.
+	for k, v := c.First(); k != nil; k, v = c.Next() {
+		if err := fn(k, v); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // page returns a reference to the page with a given id.
 // If page has been written to then a temporary bufferred page is returned.
 func (t *Transaction) page(id pgid) *page {