core: ObjectStorage.Begin and TxObjectStorage

5 files changed, 105 insertions, 8 deletions

diff --git a/storage/filesystem/object.go b/storage/filesystem/object.go
index 3888fd8..0696c70 100644
--- a/storage/filesystem/object.go
+++ b/storage/filesystem/object.go
@@ -217,3 +217,20 @@ func buildIndexFromIdxfile(fs fs.FS, path string) (index.Index, error) {
 
 	return index.NewFromIdx(f)
 }
+func (o *ObjectStorage) Begin() core.TxObjectStorage {
+	return &TxObjectStorage{}
+}
+
+type TxObjectStorage struct{}
+
+func (tx *TxObjectStorage) Set(obj core.Object) (core.Hash, error) {
+	return core.ZeroHash, fmt.Errorf("not implemented yet")
+}
+
+func (tx *TxObjectStorage) Commit() error {
+	return fmt.Errorf("not implemented yet")
+}
+
+func (tx *TxObjectStorage) Rollback() error {
+	return fmt.Errorf("not implemented yet")
+}
diff --git a/storage/filesystem/object_test.go b/storage/filesystem/object_test.go
index 361ae32..5752f79 100644
--- a/storage/filesystem/object_test.go
+++ b/storage/filesystem/object_test.go
@@ -124,9 +124,9 @@ func memStorageFromGitDir(fs fs.FS, path string) (core.ObjectStorage, error) {
 
 	sto := memory.NewStorage()
 	r := packfile.NewStream(f)
-	d := packfile.NewDecoder(r)
+	d := packfile.NewDecoder(r, sto.ObjectStorage())
 
-	err = d.Decode(sto.ObjectStorage())
+	err = d.Decode()
 	if err != nil {
 		return nil, err
 	}
@@ -267,8 +267,8 @@ func memStorageFromDirPath(fs fs.FS, path string) (core.ObjectStorage, error) {
 	}
 
 	r := packfile.NewStream(f)
-	d := packfile.NewDecoder(r)
-	err = d.Decode(sto.ObjectStorage())
+	d := packfile.NewDecoder(r, sto.ObjectStorage())
+	err = d.Decode()
 	if err != nil {
 		return nil, err
 	}
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) {
diff --git a/storage/test/storage_suite.go b/storage/test/storage_suite.go
index c68a089..b4ba306 100644
--- a/storage/test/storage_suite.go
+++ b/storage/test/storage_suite.go
@@ -114,3 +114,48 @@ func (s *BaseStorageSuite) TestObjectStorageIter(c *C) {
 		c.Assert(found, Equals, true, Commentf("Object of type %s not found", to.Type.String()))
 	}
 }
+
+func (s *BaseStorageSuite) TestTxObjectStorageSetAndCommit(c *C) {
+	tx := s.ObjectStorage.Begin()
+	for _, o := range s.testObjects {
+		h, err := tx.Set(o.Object)
+		c.Assert(err, IsNil)
+		c.Assert(h.String(), Equals, o.Hash)
+	}
+
+	iter, err := s.ObjectStorage.Iter(core.AnyObject)
+	c.Assert(err, IsNil)
+	_, err = iter.Next()
+	c.Assert(err, Equals, io.EOF)
+
+	err = tx.Commit()
+	c.Assert(err, IsNil)
+
+	iter, err = s.ObjectStorage.Iter(core.AnyObject)
+	c.Assert(err, IsNil)
+
+	var count int
+	iter.ForEach(func(o core.Object) error {
+		count++
+		return nil
+	})
+
+	c.Assert(count, Equals, 4)
+}
+
+func (s *BaseStorageSuite) TestTxObjectStorageSetAndRollback(c *C) {
+	tx := s.ObjectStorage.Begin()
+	for _, o := range s.testObjects {
+		h, err := tx.Set(o.Object)
+		c.Assert(err, IsNil)
+		c.Assert(h.String(), Equals, o.Hash)
+	}
+
+	err := tx.Rollback()
+	c.Assert(err, IsNil)
+
+	iter, err := s.ObjectStorage.Iter(core.AnyObject)
+	c.Assert(err, IsNil)
+	_, err = iter.Next()
+	c.Assert(err, Equals, io.EOF)
+}