From 56b825fb56a29da6e460475fcd5a98feb795f194 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Sat, 15 Feb 2014 14:54:45 -0700
Subject: [PATCH] Add transactional blocks.

---
 db.go           | 88 +++++++++++++++++++++----------------------------
 db_test.go      | 18 ++++++++++
 example_test.go | 27 +++++++++++++++
 3 files changed, 83 insertions(+), 50 deletions(-)

diff --git a/db.go b/db.go
index a25e41a..34de9f3 100644
--- a/db.go
+++ b/db.go
@@ -330,6 +330,26 @@ func (db *DB) removeTransaction(t *Transaction) {
 	}
 }
 
+// Do executes a function within the context of a RWTransaction.
+// If no error is returned from the function then the transaction is committed.
+// If an error is returned then the entire transaction is rolled back.
+// Any error that is returned from the function or returned from the commit is
+// returned from the Do() method.
+func (db *DB) Do(fn func(*RWTransaction) error) error {
+	t, err := db.RWTransaction()
+	if err != nil {
+		return err
+	}
+
+	// If an error is returned from the function then rollback and return error.
+	if err := fn(t); err != nil {
+		t.Rollback()
+		return err
+	}
+
+	return t.Commit()
+}
+
 // 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) {
@@ -355,51 +375,31 @@ func (db *DB) Buckets() ([]*Bucket, error) {
 // This function can return an error if the bucket already exists, if the name
 // is blank, or the bucket name is too long.
 func (db *DB) CreateBucket(name string) error {
-	t, err := db.RWTransaction()
-	if err != nil {
-		return err
-	}
-
-	if err := t.CreateBucket(name); err != nil {
-		t.Rollback()
-		return err
-	}
-
-	return t.Commit()
+	return db.Do(func(t *RWTransaction) error {
+		return t.CreateBucket(name)
+	})
 }
 
 // DeleteBucket removes a bucket from the database.
 // Returns an error if the bucket does not exist.
 func (db *DB) DeleteBucket(name string) error {
-	t, err := db.RWTransaction()
-	if err != nil {
-		return err
-	}
-
-	if err := t.DeleteBucket(name); err != nil {
-		t.Rollback()
-		return err
-	}
-
-	return t.Commit()
+	return db.Do(func(t *RWTransaction) error {
+		return t.DeleteBucket(name)
+	})
 }
 
 // NextSequence returns an autoincrementing integer for the bucket.
 // This function can return an error if the bucket does not exist.
 func (db *DB) NextSequence(name string) (int, error) {
-	t, err := db.RWTransaction()
+	var seq int 
+	err := db.Do(func(t *RWTransaction) error {
+		var err error
+		seq, err = t.NextSequence(name)
+		return err
+	})
 	if err != nil {
 		return 0, err
 	}
-
-	seq, err := t.NextSequence(name)
-	if err != nil {
-		t.Rollback()
-		return 0, err
-	}
-	if err := t.Commit(); err != nil {
-		return 0, err
-	}
 	return seq, nil
 }
 
@@ -417,29 +417,17 @@ func (db *DB) Get(name string, key []byte) ([]byte, error) {
 // Put sets the value for a key in a bucket.
 // Returns an error if the bucket is not found, if key is blank, if the key is too large, or if the value is too large.
 func (db *DB) Put(name string, key []byte, value []byte) error {
-	t, err := db.RWTransaction()
-	if err != nil {
-		return err
-	}
-	if err := t.Put(name, key, value); err != nil {
-		t.Rollback()
-		return err
-	}
-	return t.Commit()
+	return db.Do(func(t *RWTransaction) error {
+		return t.Put(name, key, value)
+	})
 }
 
 // Delete removes a key from a bucket.
 // Returns an error if the bucket cannot be found.
 func (db *DB) Delete(name string, key []byte) error {
-	t, err := db.RWTransaction()
-	if err != nil {
-		return err
-	}
-	if err := t.Delete(name, key); err != nil {
-		t.Rollback()
-		return err
-	}
-	return t.Commit()
+	return db.Do(func(t *RWTransaction) error {
+		return t.Delete(name, key)
+	})
 }
 
 // Copy writes the entire database to a writer.
diff --git a/db_test.go b/db_test.go
index 2682f55..1a9aa02 100644
--- a/db_test.go
+++ b/db_test.go
@@ -178,6 +178,24 @@ func TestDBDelete(t *testing.T) {
 	})
 }
 
+// Ensure a database can provide a transactional block.
+func TestDBTransactionBlock(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		err := db.Do(func(txn *RWTransaction) error {
+			txn.CreateBucket("widgets")
+			txn.Put("widgets", []byte("foo"), []byte("bar"))
+			txn.Put("widgets", []byte("baz"), []byte("bat"))
+			txn.Delete("widgets", []byte("foo"))
+			return nil
+		})
+		assert.NoError(t, err)
+		value, _ := db.Get("widgets", []byte("foo"))
+		assert.Nil(t, value)
+		value, _ = db.Get("widgets", []byte("baz"))
+		assert.Equal(t, value, []byte("bat"))
+	})
+}
+
 // Ensure that the database can be copied to a writer.
 func TestDBCopy(t *testing.T) {
 	t.Skip("pending") // TODO(benbjohnson)
diff --git a/example_test.go b/example_test.go
index 892807e..655e283 100644
--- a/example_test.go
+++ b/example_test.go
@@ -60,6 +60,33 @@ func ExampleDB_Delete() {
 	// The value of 'foo' is now: nil
 }
 
+func ExampleDB_Do() {
+	// Open the database.
+	var db DB
+	db.Open("/tmp/bolt/db_do.db", 0666)
+	defer db.Close()
+
+	// Execute several commands within a write transaction.
+	err := db.Do(func(t *RWTransaction) error {
+		if err := t.CreateBucket("widgets"); err != nil {
+			return err
+		}
+		if err := t.Put("widgets", []byte("foo"), []byte("bar")); err != nil {
+			return err
+		}
+		return nil
+	})
+
+	// If our transactional block didn't return an error then our data is saved.
+	if err == nil {
+		value, _ := db.Get("widgets", []byte("foo"))
+		fmt.Printf("The value of 'foo' is: %s\n", string(value))
+	}
+
+	// Output:
+	// The value of 'foo' is: bar
+}
+
 func ExampleRWTransaction() {
 	// Open the database.
 	var db DB