def end(self, handle_invalidations=None):
write(self.s, 'C')
n = read_int4(self.s)
oid_list = []
if n != 0:
packed_oids = read(self.s, n * 8)
oid_list = split_oids(packed_oids)
try:
handle_invalidations(oid_list)
except ConflictError:
self.transaction_new_oids.reverse()
self.oid_pool.extend(self.transaction_new_oids)
assert len(self.oid_pool) == len(set(self.oid_pool))
self.begin() # clear out records and transaction_new_oids.
write_int4(self.s, 0) # Tell server we are done.
raise
tdata = []
for oid, record in iteritems(self.records):
tdata.append(int4_to_str(8 + len(record)))
tdata.append(as_bytes(oid))
tdata.append(record)
tdata = join_bytes(tdata)
write_int4_str(self.s, tdata)
self.records.clear()
if len(tdata) > 0:
status = read(self.s, 1)
if status == STATUS_OKAY:
pass
elif status == STATUS_INVALID:
raise WriteConflictError()
else:
raise ProtocolError('server returned invalid status %r' %
status)
def end(self, handle_invalidations=None):
write(self.s, 'C')
n = read_int4(self.s)
oid_list = []
if n != 0:
packed_oids = read(self.s, n*8)
oid_list = split_oids(packed_oids)
try:
handle_invalidations(oid_list)
except ConflictError:
self.transaction_new_oids.reverse()
self.oid_pool.extend(self.transaction_new_oids)
assert len(self.oid_pool) == len(set(self.oid_pool))
self.begin() # clear out records and transaction_new_oids.
write_int4(self.s, 0) # Tell server we are done.
raise
tdata = []
for oid, record in iteritems(self.records):
tdata.append(int4_to_str(8 + len(record)))
tdata.append(as_bytes(oid))
tdata.append(record)
tdata = join_bytes(tdata)
write_int4_str(self.s, tdata)
self.records.clear()
if len(tdata) > 0:
status = read(self.s, 1)
if status == STATUS_OKAY:
pass
elif status == STATUS_INVALID:
raise WriteConflictError()
else:
raise ProtocolError('server returned invalid status %r' % status)
def check_receive(self):
class Dribble:
def recv(x, n):
return as_bytes(choice(['a', 'bb'])[:n])
fake_socket = Dribble()
read(fake_socket, 30)
def f(self):
class FakeSocket(object):
def recv(self, n):
if n > 10:
return as_bytes('')
return as_bytes('x')
def send(self, s):
return len(s)
s = FakeSocket()
write(s, 'x' * 2000000)
read(s, 8)
raises(ShortRead, read, s, 11)
def _get_load_response(self, oid):
status = read(self.s, 1)
if status == STATUS_OKAY:
pass
elif status == STATUS_INVALID:
raise ReadConflictError([oid])
elif status == STATUS_KEYERROR:
raise DurusKeyError(oid)
else:
raise ProtocolError('status=%r, oid=%r' % (status, oid))
n = read_int4(self.s)
record = read(self.s, n)
return record
def _get_load_response(self, oid):
status = read(self.s, 1)
if status == STATUS_OKAY:
pass
elif status == STATUS_INVALID:
raise ReadConflictError([oid])
elif status == STATUS_KEYERROR:
raise DurusKeyError(oid)
else:
raise ProtocolError('status=%r, oid=%r' % (status, oid))
n = read_int4(self.s)
record = read(self.s, n)
return record
def server_protocol(self):
"""Get the protocol used by the server.
:rtype: 4-byte string
"""
write(self.socket, 'V')
return read(self.socket, 4)
def _enumerate_database_names(self):
count = read_int4(self.socket)
while count > 0:
count -= 1
length = read_int4(self.socket)
database_name = read(self.socket, length)
yield database_name
def handle_B(self, s):
# bulk read of objects
number_of_oids = read_int4(s)
oid_str = read(s, 8 * number_of_oids)
oids = split_oids(oid_str)
for oid in oids:
self._send_load_response(s, oid)
def has_format(klass, file):
file.seek(0)
try:
prefix = read(file, len(klass.prefix))
except ShortRead:
return False
return klass.prefix == prefix
def has_format(klass, file):
file.seek(0)
try:
prefix = read(file, len(klass.prefix))
except ShortRead:
return False
return klass.prefix == prefix
def handle_B(self, s):
# bulk read of objects
number_of_oids = read_int4(s)
oid_str = read(s, 8 * number_of_oids)
oids = split_oids(oid_str)
for oid in oids:
self._send_load_response(s, oid)
def has_format(klass, file):
file.seek(0)
try:
if klass.MAGIC == read(file, len(klass.MAGIC)):
return True
except ShortRead:
pass
return False
def sync(self):
self._send_command('S')
n = read_int4(self.socket)
if n == 0:
packed_oids = ''
else:
packed_oids = read(self.socket, 8 * n)
return split_oids(packed_oids)
def sync(self):
write(self.s, 'S')
n = read_int4(self.s)
if n == 0:
packed_oids = ''
else:
packed_oids = read(self.s, n * 8)
return split_oids(packed_oids)
def handle_V(self, s):
# Verify protocol version match.
client_protocol = read(s, 4)
log(10, 'Client Protocol: %s', str_to_int4(client_protocol))
assert len(self.protocol) == 4
write(s, self.protocol)
if client_protocol != self.protocol:
raise ClientError("Protocol not supported.")
def b(self):
s = BytesIO()
for x in ("", "a", "ab", "a" * 1000):
x = as_bytes(x)
s.seek(0)
write(s, x)
s.seek(0)
assert x == read(s, len(x))
def sync(self):
write(self.s, 'S')
n = read_int4(self.s)
if n == 0:
packed_oids = ''
else:
packed_oids = read(self.s, n*8)
return split_oids(packed_oids)
def b(self):
s = BytesIO()
for x in ('', 'a', 'ab', 'a' * 1000):
x = as_bytes(x)
s.seek(0)
write(s, x)
s.seek(0)
assert x == read(s, len(x))
def handle_V(self, s):
# Verify protocol version match.
client_protocol = read(s, 4)
log(10, 'Client Protocol: %s', str_to_int4(client_protocol))
assert len(self.protocol) == 4
write(s, self.protocol)
if client_protocol != self.protocol:
raise ClientError("Protocol not supported.")
def generate_shelf(klass, file, items):
"""(File, [(str, str)])
This returns a generator that writes a new Shelf into file,
iterating once through the given items.
The use of an iterator makes it possible to build a new Shelf
incrementally.
"""
file.seek_end()
if not file.tell() == 0:
raise ValueError("Expected %s to be empty." % file)
write(file, klass.prefix)
if not items:
# Just write an empty transaction.
write_int8(file, 0)
# Write an empty index array.
offset_map = OffsetMap(file)
else:
# Write a transaction here with the given items.
transaction_start = file.tell()
# Write a placeholder for the length.
write_int8(file, 0)
# Loop over the items, writing their records.
# Keep track of max_key and max_offset.
max_key = 0
max_offset = 0
n = 0
for name, value in items:
max_key = max(max_key, str_to_int8(name))
max_offset = max(max_offset, file.tell())
write_int8(file, len(name) + len(value))
write(file, name)
write(file, value)
n += 1
yield n
transaction_end = file.tell()
# Write the correct transaction length.
file.seek(transaction_start)
write_int8(file, transaction_end - transaction_start - 8)
# Write the empty array with the calculated dimensions.
file.seek(transaction_end)
for step in OffsetMap.generate(file, max_key, max_offset):
yield step
offset_map = OffsetMap(file)
# Now read through the records and record the offsets in the array.
file.seek(transaction_start + 8)
while file.tell() < transaction_end:
position = file.tell()
record_length = read_int8(file)
name = read(file, 8)
k = str_to_int8(name)
offset_map[k] = position
file.seek(position + 8 + record_length)
n -= 1
yield n
for index in offset_map.gen_stitch():
yield index
def generate_shelf(klass, file, items):
"""(File, [(str, str)])
This returns a generator that writes a new Shelf into file,
iterating once through the given items.
The use of an iterator makes it possible to build a new Shelf
incrementally.
"""
file.seek_end()
if not file.tell() == 0:
raise ValueError("Expected %s to be empty." % file)
write(file, klass.prefix)
if not items:
# Just write an empty transaction.
write_int8(file, 0)
# Write an empty index array.
offset_map = OffsetMap(file)
else:
# Write a transaction here with the given items.
transaction_start = file.tell()
# Write a placeholder for the length.
write_int8(file, 0)
# Loop over the items, writing their records.
# Keep track of max_key and max_offset.
max_key = 0
max_offset = 0
n = 0
for name, value in items:
max_key = max(max_key, str_to_int8(name))
max_offset = max(max_offset, file.tell())
write_int8(file, len(name) + len(value))
write(file, name)
write(file, value)
n += 1
yield n
transaction_end = file.tell()
# Write the correct transaction length.
file.seek(transaction_start)
write_int8(file, transaction_end - transaction_start - 8)
# Write the empty array with the calculated dimensions.
file.seek(transaction_end)
for step in OffsetMap.generate(file, max_key, max_offset):
yield step
offset_map = OffsetMap(file)
# Now read through the records and record the offsets in the array.
file.seek(transaction_start + 8)
while file.tell() < transaction_end:
position = file.tell()
record_length = read_int8(file)
name = read(file, 8)
k = str_to_int8(name)
offset_map[k] = position
file.seek(position + 8 + record_length)
n -= 1
yield n
for index in offset_map.gen_stitch():
yield index
def handle(self, s):
command_byte = read(s, 1)[0]
if type(command_byte) is int:
command_code = chr(command_byte)
else:
command_code = command_byte
handler = getattr(self, 'handle_%s' % command_code, None)
if handler is None:
raise ClientError('No such command code: %r' % command_code)
handler(s)
def handle(self, s):
command_byte = read(s, 1)[0]
if type(command_byte) is int:
command_code = chr(command_byte)
else:
command_code = command_byte
handler = getattr(self, 'handle_%s' % command_code, None)
if handler is None:
raise ClientError('No such command code: %r' % command_code)
handler(s)
def new_oid(self):
if not self.oid_pool:
batch = self.oid_pool_size
write(self.s, 'M%s' % chr(batch))
self.oid_pool = split_oids(read(self.s, 8 * batch))
self.oid_pool.reverse()
assert len(self.oid_pool) == len(set(self.oid_pool))
oid = self.oid_pool.pop()
assert oid not in self.oid_pool
self.transaction_new_oids.append(oid)
return oid
def new_oid(self):
if not self.oid_pool:
batch = self.oid_pool_size
write(self.s, 'M%s' % chr(batch))
self.oid_pool = split_oids(read(self.s, 8 * batch))
self.oid_pool.reverse()
assert len(self.oid_pool) == len(set(self.oid_pool))
oid = self.oid_pool.pop()
assert oid not in self.oid_pool
self.transaction_new_oids.append(oid)
return oid
def new_oid(self):
if not self.oid_pool:
batch = self.oid_pool_size
self._send_command('M')
write(self.socket, chr(batch))
self.oid_pool = split_oids(read(self.socket, 8 * batch))
self.oid_pool.reverse()
assert len(self.oid_pool) == len(set(self.oid_pool))
oid = self.oid_pool.pop()
assert oid not in self.oid_pool
self.transaction_new_oids.append(oid)
return oid
def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, address=None):
self.address = SocketAddress.new(address or (host, port))
self.s = self.address.get_connected_socket()
assert self.s, "Could not connect to %s" % self.address
self.oid_pool = []
self.oid_pool_size = 32
self.begin()
protocol = StorageServer.protocol
assert len(protocol) == 4
write_all(self.s, 'V', protocol)
server_protocol = read(self.s, 4)
if server_protocol != protocol:
raise ProtocolError("Protocol version mismatch.")
def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, address=None):
self.address = SocketAddress.new(address or (host, port))
self.s = self.address.get_connected_socket()
assert self.s, "Could not connect to %s" % self.address
self.oid_pool = []
self.oid_pool_size = 32
self.begin()
protocol = StorageServer.protocol
assert len(protocol) == 4
write_all(self.s, 'V', protocol)
server_protocol = read(self.s, 4)
if server_protocol != protocol:
raise ProtocolError("Protocol version mismatch.")
def _build_index(self, repair):
self.fp.seek(0)
if read(self.fp, len(self.MAGIC)) != self.MAGIC:
raise IOError("invalid storage (missing magic in %r)" % self.fp)
index_offset = read_int8(self.fp)
assert index_offset > 0
self.fp.seek(index_offset)
tmp_index = loads(decompress(read_int8_str(self.fp)))
self.index = {}
def oid_as_bytes(oid):
if isinstance(oid, byte_string):
return oid
else:
return oid.encode('latin1')
for tmp_oid in tmp_index:
self.index[oid_as_bytes(tmp_oid)] = tmp_index[tmp_oid]
del tmp_index
while 1:
# Read one transaction each time here.
oids = {}
transaction_offset = self.fp.tell()
try:
while 1:
object_record_offset = self.fp.tell()
record = self._read_block()
if len(record) == 0:
break # normal termination
if len(record) < 12:
raise ShortRead("Bad record size")
oid = record[0:8]
oids[oid] = object_record_offset
# We've reached the normal end of a transaction.
self.index.update(oids)
oids.clear()
except (ValueError, IOError):
if self.fp.tell() > transaction_offset:
if not repair:
raise
# The transaction was malformed. Attempt repair.
self.fp.seek(transaction_offset)
self.fp.truncate()
break
def read_transaction_offsets(file, repair=False):
"""
Read the offsets of one (Shelf-format) transaction in the file.
If repair is True and the file ends in something other than a well
formed transaction, the file is truncated to remove the ugly
ending. This ugliness might occur if the power fails in the middle
of writing a transaction.
"""
transaction_start = transaction_end = position = file.tell()
try:
transaction_length = read_int8(file)
transaction_end = transaction_start + 8 + transaction_length
transaction_offsets = {}
while file.tell() < transaction_end:
position = file.tell()
record_length = read_int8(file)
identifier = read(file, 8)
transaction_offsets[identifier] = position
file.seek(position + 8 + record_length)
if file.tell() != transaction_end:
raise ShortRead
return transaction_offsets
except ShortRead:
position = file.tell()
if position > transaction_start:
if repair:
file.seek(transaction_start)
file.truncate()
else:
e = sys.exc_info()[1]
e.args = repr(
dict(transaction_start=transaction_start,
transaction_end=transaction_end,
position=position))
raise
return None
def read_transaction_offsets(file, repair=False):
"""
Read the offsets of one (Shelf-format) transaction in the file.
If repair is True and the file ends in something other than a well
formed transaction, the file is truncated to remove the ugly
ending. This ugliness might occur if the power fails in the middle
of writing a transaction.
"""
transaction_start = transaction_end = position = file.tell()
try:
transaction_length = read_int8(file)
transaction_end = transaction_start + 8 + transaction_length
transaction_offsets = {}
while file.tell() < transaction_end:
position = file.tell()
record_length = read_int8(file)
identifier = read(file, 8)
transaction_offsets[identifier] = position
file.seek(position + 8 + record_length)
if file.tell() != transaction_end:
raise ShortRead
return transaction_offsets
except ShortRead:
position = file.tell()
if position > transaction_start:
if repair:
file.seek(transaction_start)
file.truncate()
else:
e = sys.exc_info()[1]
e.args = repr(dict(
transaction_start=transaction_start,
transaction_end = transaction_end,
position=position))
raise
return None
def pack(self):
write(self.s, 'P')
status = read(self.s, 1)
if status != STATUS_OKAY:
raise ProtocolError('server returned invalid status %r' % status)
def check_receive(self):
class Dribble:
def recv(x, n):
return as_bytes(choice(['a', 'bb'])[:n])
fake_socket = Dribble()
read(fake_socket, 30)
def handle_M(self, s):
# new OIDs
count = ord(read(s, 1))
log(10, "oids: %s", count)
write(s, join_bytes(self._new_oids(s, count)))
def handle_L(self, s):
# load
oid = read(s, 8)
self._send_load_response(s, oid)
def pack(self):
write(self.s, 'P')
status = read(self.s, 1)
if status != STATUS_OKAY:
raise ProtocolError('server returned invalid status %r' % status)
def pack(self):
self._send_command('P')
status = read(self.socket, 1)
if status != STATUS_OKAY:
raise ProtocolError('server returned invalid status %r' % status)
def handle_L(self, s):
# load
oid = read(s, 8)
self._send_load_response(s, oid)
def handle_M(self, s):
# new OIDs
count = ord(read(s, 1))
log(10, "oids: %s", count)
write(s, join_bytes(self._new_oids(s, count)))
