def SaveTemporaryFile(fp):
"""Store incoming database file in a temporary directory."""
loc = data_store.DB.Location()
if not os.path.exists(loc):
return False
if not os.path.isdir(loc):
return False
# Read DataServerFileCopy object.
filecopy_len_str = fp.read(sutils.SIZE_PACKER.size)
filecopy_len = sutils.SIZE_PACKER.unpack(filecopy_len_str)[0]
filecopy = rdf_data_server.DataServerFileCopy(fp.read(filecopy_len))
rebdir = _CreateDirectory(loc, filecopy.rebalance_id)
filedir = utils.JoinPath(rebdir, filecopy.directory)
try:
os.makedirs(filedir)
except OSError:
pass
filepath = utils.JoinPath(filedir, filecopy.filename)
logging.info("Writing to file %s", filepath)
with open(filepath, "wb") as wp:
# We need to uncompress the file stream.
decompressor = zlib.decompressobj()
while True:
block_len_str = fp.read(sutils.SIZE_PACKER.size)
block_len = sutils.SIZE_PACKER.unpack(block_len_str)[0]
if not block_len:
break
block = fp.read(block_len)
to_decompress = decompressor.unconsumed_tail + block
while to_decompress:
decompressed = decompressor.decompress(to_decompress)
if decompressed:
wp.write(decompressed)
to_decompress = decompressor.unconsumed_tail
else:
to_decompress = ""
# Deal with remaining data.
remaining = decompressor.flush()
if remaining:
wp.write(remaining)
if os.path.getsize(filepath) != filecopy.size:
logging.error("Size of file %s is not %d", filepath, filecopy.size)
return False
return True
def __init__(self, rebalance, directory, filename, fullpath):
filesize = os.path.getsize(fullpath)
filecopy = rdf_data_server.DataServerFileCopy(rebalance_id=rebalance.id,
directory=directory,
filename=filename,
size=filesize)
filecopy_str = filecopy.SerializeToString()
self.header = sutils.SIZE_PACKER.pack(len(filecopy_str))
self.header += filecopy_str
self.header = StringIO.StringIO(self.header)
self.fp = open(fullpath, "rb")
self.compressor = zlib.compressobj(COMPRESSION_LEVEL)
# Buffered compressed data that needs to be read.
self.buffered = ""
# Flag to mark end of database file.
self.end_of_file = False
# Flag to mark if we can no longer use read().
self.end_of_stream = False
# Flag to indicate that we are going to read the header first.
self.read_header = True