MemoryDB: Implemented direct stream copy, allowing to flush memory db content into any other object db for permanent storage

Author: Byron

db/loose.py

@@ -13,6 +13,7 @@
 from gitdb.stream import (
 		DecompressMemMapReader,
 		FDCompressedSha1Writer,
+		FDStream,
 		Sha1Writer
 	)
 
@@ -43,6 +44,7 @@
 
 import tempfile
 import mmap
+import sys
 import os
 
 
@@ -153,13 +155,20 @@ def store(self, istream):
 		if writer is None:
 			# open a tmp file to write the data to
 			fd, tmp_path = tempfile.mkstemp(prefix='obj', dir=self._root_path)
-			writer = FDCompressedSha1Writer(fd)
+			
+			if istream.sha is None:
+				writer = FDCompressedSha1Writer(fd)
+			else:
+				writer = FDStream(fd)
+			# END handle direct stream copies
 		# END handle custom writer
 
 		try:
 			try:
 				if istream.sha is not None:
-					stream_copy(istream.read, writer.write, istream.size, self.stream_chunk_size)
+					# copy as much as possible, the actual uncompressed item size might
+					# be smaller than the compressed version
+					stream_copy(istream.read, writer.write, sys.maxint, self.stream_chunk_size)
 				else:
 					# write object with header, we have to make a new one
 					write_object(istream.type, istream.size, istream.read, writer.write,
@@ -175,18 +184,23 @@ def store(self, istream):
 				writer.close()
 		# END assure target stream is closed
 
-		sha = istream.sha or writer.sha(as_hex=True)
+		hexsha = None
+		if istream.sha:
+			hexsha = istream.hexsha
+		else:
+			hexsha = writer.sha(as_hex=True)
+		# END handle sha
 
 		if tmp_path:
-			obj_path = self.db_path(self.object_path(sha))
+			obj_path = self.db_path(self.object_path(hexsha))
 			obj_dir = dirname(obj_path)
 			if not isdir(obj_dir):
 				mkdir(obj_dir)
 			# END handle destination directory
 			rename(tmp_path, obj_path)
 		# END handle dry_run
 
-		istream.sha = sha
+		istream.sha = hexsha
 		return istream
 
 	def sha_iter(self):

db/mem.py

@@ -5,7 +5,11 @@
 						ObjectDBW
 					)
 
-from gitdb.base import OStream
+from gitdb.base import (
+							OStream,
+							IStream,
+						)
+
 from gitdb.util import to_bin_sha
 from gitdb.exc import (
 						BadObject,
@@ -16,6 +20,8 @@
 							DecompressMemMapReader,
 						)
 
+from cStringIO import StringIO
+
 __all__ = ("MemoryDB", )
 
 class MemoryDB(ObjectDBR, ObjectDBW):
@@ -78,3 +84,28 @@ def size(self):
 
 	def sha_iter(self):
 		return self._cache.iterkeys()
+		
+		
+	#{ Interface 
+	def stream_copy(self, sha_iter, odb):
+		"""Copy the streams as identified by sha's yielded by sha_iter into the given odb
+		The streams will be copied directly
+		:note: the object will only be written if it did not exist in the target db
+		:return: amount of streams actually copied into odb. If smaller than the amount
+			of input shas, one or more objects did already exist in odb"""
+		count = 0
+		for sha in sha_iter:
+			if odb.has_object(sha):
+				continue
+			# END check object existance
+			
+			ostream = self.stream(sha)
+			# compressed data including header
+			sio = StringIO(ostream.stream.data())
+			istream = IStream(ostream.type, ostream.size, sio, sha)
+			
+			odb.store(istream)
+			count += 1
+		# END for each sha
+		return count
+	#} END interface

stream.py

@@ -25,21 +25,6 @@
 
 #{ RO Streams
 
-class NullStream(object):
-	"""A stream that does nothing but providing a stream interface.
-	Use it like /dev/null"""
-	__slots__ = tuple()
-		
-	def read(self, size=0):
-		return ''
-		
-	def close(self):
-		pass
-		
-	def write(self, data):
-		return len(data)
-
-
 class DecompressMemMapReader(LazyMixin):
 	"""Reads data in chunks from a memory map and decompresses it. The client sees 
 	only the uncompressed data, respective file-like read calls are handling on-demand
@@ -113,6 +98,8 @@ def _parse_header_info(self):
 
 		return type, size
 
+	#{ Interface 
+	
 	@classmethod
 	def new(self, m, close_on_deletion=False):
 		"""Create a new DecompressMemMapReader instance for acting as a read-only stream
@@ -125,6 +112,10 @@ def new(self, m, close_on_deletion=False):
 		type, size = inst._parse_header_info()
 		return type, size, inst
 
+	def data(self):
+		""":return: random access compatible data we are working on"""
+		return self._m 
+	
 	def compressed_bytes_read(self):
 		""":return: number of compressed bytes read. This includes the bytes it 
 		took to decompress the header ( if there was one )"""
@@ -171,6 +162,8 @@ def compressed_bytes_read(self):
 		# from the count already
 		return self._cbr
 
+	#} END interface 
+		
 	def seek(self, offset, whence=os.SEEK_SET):
 		"""Allows to reset the stream to restart reading
 		:raise ValueError: If offset and whence are not 0"""
@@ -567,4 +560,32 @@ def close(self):
 
 	#} END stream interface
 
+class FDStream(object):
+	"""Simple wrapper around a file descriptor"""
+	__slots__ = "_fd"
+	def __init__(self, fd):
+		self._fd = fd
+		
+	def write(self, data):
+		return write(self._fd, data)
+		
+	def close(self):
+		close(self._fd)
+	
+	
+
+class NullStream(object):
+	"""A stream that does nothing but providing a stream interface.
+	Use it like /dev/null"""
+	__slots__ = tuple()
+		
+	def read(self, size=0):
+		return ''
+		
+	def close(self):
+		pass
+		
+	def write(self, data):
+		return len(data)
+
 #} END W streams

test/db/test_mem.py

@@ -1,10 +1,26 @@
 from lib import *
-from gitdb.db import MemoryDB
+from gitdb.db import (
+						MemoryDB,
+						LooseObjectDB
+					)
 
 class TestMemoryDB(TestDBBase):
 
-	def test_writing(self):
+	@with_rw_directory
+	def test_writing(self, path):
 		mdb = MemoryDB()
 
 		# write data
 		self._assert_object_writing_simple(mdb)
+		
+		# test stream copy
+		ldb = LooseObjectDB(path)
+		assert ldb.size() == 0
+		num_streams_copied = mdb.stream_copy(mdb.sha_iter(), ldb)
+		assert num_streams_copied == mdb.size()
+		
+		assert ldb.size() == mdb.size()
+		for sha in mdb.sha_iter():
+			assert ldb.has_object(sha)
+			assert ldb.stream(sha).read() == mdb.stream(sha).read() 
+		# END verify objects where copied and are equal

test/test_stream.py

@@ -50,7 +50,7 @@ def _assert_stream_reader(self, stream, cdata, rewind_stream=lambda s: None):
 		# END handle rest
 
 		if isinstance(stream, DecompressMemMapReader):
-			assert len(stream._m) == stream.compressed_bytes_read()
+			assert len(stream.data()) == stream.compressed_bytes_read()
 		# END handle special type
 
 		rewind_stream(stream)
@@ -60,7 +60,7 @@ def _assert_stream_reader(self, stream, cdata, rewind_stream=lambda s: None):
 		assert rdata == cdata
 
 		if isinstance(stream, DecompressMemMapReader):
-			assert len(stream._m) == stream.compressed_bytes_read()
+			assert len(stream.data()) == stream.compressed_bytes_read()
 		# END handle special type
 
 	def test_decompress_reader(self):