core: ObjectStorage.Begin and TxObjectStorage

2 files changed, 39 insertions, 4 deletions

diff --git a/storage/memory/storage.go b/storage/memory/storage.go
index c8a152d..01d5fa5 100644
--- a/storage/memory/storage.go
+++ b/storage/memory/storage.go
@@ -166,6 +166,41 @@ func flattenObjectMap(m map[core.Hash]core.Object) []core.Object {
 	return objects
 }
 
+func (o *ObjectStorage) Begin() core.TxObjectStorage {
+	return &TxObjectStorage{
+		Storage: o,
+		Objects: make(map[core.Hash]core.Object, 0),
+	}
+}
+
+type TxObjectStorage struct {
+	Storage *ObjectStorage
+	Objects map[core.Hash]core.Object
+}
+
+func (tx *TxObjectStorage) Set(obj core.Object) (core.Hash, error) {
+	h := obj.Hash()
+	tx.Objects[h] = obj
+
+	return h, nil
+}
+
+func (tx *TxObjectStorage) Commit() error {
+	for h, obj := range tx.Objects {
+		delete(tx.Objects, h)
+		if _, err := tx.Storage.Set(obj); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (tx *TxObjectStorage) Rollback() error {
+	tx.Objects = make(map[core.Hash]core.Object, 0)
+	return nil
+}
+
 type ReferenceStorage map[core.ReferenceName]*core.Reference
 
 // Set stores a reference.
diff --git a/storage/memory/storage_test.go b/storage/memory/storage_test.go
index 61a810e..2c90738 100644
--- a/storage/memory/storage_test.go
+++ b/storage/memory/storage_test.go
@@ -6,19 +6,19 @@ import (
 
 	. "gopkg.in/check.v1"
 	"gopkg.in/src-d/go-git.v4/core"
-	. "gopkg.in/src-d/go-git.v4/storage/test"
+	"gopkg.in/src-d/go-git.v4/storage/test"
 )
 
 func Test(t *testing.T) { TestingT(t) }
 
 type StorageSuite struct {
-	BaseStorageSuite
+	test.BaseStorageSuite
 }
 
 var _ = Suite(&StorageSuite{})
 
-func (s *StorageSuite) SetUpSuite(c *C) {
-	s.BaseStorageSuite = NewBaseStorageSuite(NewStorage().ObjectStorage())
+func (s *StorageSuite) SetUpTest(c *C) {
+	s.BaseStorageSuite = test.NewBaseStorageSuite(NewStorage().ObjectStorage())
 }
 
 func (s *StorageSuite) TestStorageObjectStorage(c *C) {