Added CommitBatch to the versionDB

database/versiondb/versiondb.go

@@ -184,29 +184,55 @@ func (db *Database) Commit() error {
 	db.lock.Lock()
 	defer db.lock.Unlock()
 
-	if db.mem == nil {
-		return database.ErrClosed
+	batch, err := db.commitBatch()
+	if err != nil {
+		return err
 	}
-	if len(db.mem) == 0 {
-		return nil
+	if err := batch.Write(); err != nil {
+		return err
+	}
+	db.abort()
+	return nil
+}
+
+// Abort all changes to the underlying database
+func (db *Database) Abort() {
+	db.lock.Lock()
+	defer db.lock.Unlock()
+
+	db.abort()
+}
+
+func (db *Database) abort() { db.mem = make(map[string]valueDelete, memdb.DefaultSize) }
+
+// CommitBatch returns a batch that will commit all pending writes to the underlying database
+func (db *Database) CommitBatch() (database.Batch, error) {
+	db.lock.Lock()
+	defer db.lock.Unlock()
+
+	return db.commitBatch()
+}
+
+func (db *Database) commitBatch() (database.Batch, error) {
+	if db.mem == nil {
+		return nil, database.ErrClosed
 	}
 
 	batch := db.db.NewBatch()
 	for key, value := range db.mem {
 		if value.delete {
 			if err := batch.Delete([]byte(key)); err != nil {
-				return err
+				return nil, err
 			}
 		} else if err := batch.Put([]byte(key), value.value); err != nil {
-			return err
+			return nil, err
 		}
 	}
 	if err := batch.Write(); err != nil {
-		return err
+		return nil, err
 	}
 
-	db.mem = make(map[string]valueDelete, memdb.DefaultSize)
-	return nil
+	return batch, nil
 }
 
 // Close implements the database.Database interface

database/versiondb/versiondb_test.go

@@ -256,6 +256,72 @@ func TestCommitClosedDelete(t *testing.T) {
 	}
 }
 
+func TestAbort(t *testing.T) {
+	baseDB := memdb.New()
+	db := New(baseDB)
+
+	key1 := []byte("hello1")
+	value1 := []byte("world1")
+
+	if err := db.Put(key1, value1); err != nil {
+		t.Fatalf("Unexpected error on db.Put: %s", err)
+	}
+
+	if value, err := db.Get(key1); err != nil {
+		t.Fatalf("Unexpected error on db.Get: %s", err)
+	} else if !bytes.Equal(value, value1) {
+		t.Fatalf("db.Get Returned: 0x%x ; Expected: 0x%x", value, value1)
+	} else if has, err := baseDB.Has(key1); err != nil {
+		t.Fatalf("Unexpected error on db.Has: %s", err)
+	} else if has {
+		t.Fatalf("db.Has Returned: %v ; Expected: %v", has, false)
+	}
+
+	db.Abort()
+
+	if has, err := db.Has(key1); err != nil {
+		t.Fatalf("Unexpected error on db.Has: %s", err)
+	} else if has {
+		t.Fatalf("db.Has Returned: %v ; Expected: %v", has, false)
+	} else if has, err := baseDB.Has(key1); err != nil {
+		t.Fatalf("Unexpected error on db.Has: %s", err)
+	} else if has {
+		t.Fatalf("db.Has Returned: %v ; Expected: %v", has, false)
+	}
+}
+
+func TestCommitBatch(t *testing.T) {
+	baseDB := memdb.New()
+	db := New(baseDB)
+
+	key1 := []byte("hello1")
+	value1 := []byte("world1")
+
+	if err := db.Put(key1, value1); err != nil {
+		t.Fatalf("Unexpected error on db.Put: %s", err)
+	}
+
+	batch, err := db.CommitBatch()
+	if err != nil {
+		t.Fatalf("Unexpected error on db.CommitBatch: %s", err)
+	}
+	db.Abort()
+
+	if err := batch.Write(); err != nil {
+		t.Fatalf("Unexpected error on batch.Write: %s", err)
+	}
+
+	if value, err := db.Get(key1); err != nil {
+		t.Fatalf("Unexpected error on db.Get: %s", err)
+	} else if !bytes.Equal(value, value1) {
+		t.Fatalf("db.Get Returned: 0x%x ; Expected: 0x%x", value, value1)
+	} else if value, err := baseDB.Get(key1); err != nil {
+		t.Fatalf("Unexpected error on db.Get: %s", err)
+	} else if !bytes.Equal(value, value1) {
+		t.Fatalf("db.Get Returned: 0x%x ; Expected: 0x%x", value, value1)
+	}
+}
+
 func TestSetDatabase(t *testing.T) {
 	baseDB := memdb.New()
 	newDB := memdb.New()