def test_write_2(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_4)
proto_table_storage.write(data={'test': self.EXAMPLE_PROTO_1})
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_2})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_2)
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_3})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_3)
result = proto_table_storage.read_all()
self.assertDictEqual(
result, {
'test': ProtoUtil.message_to_any(self.EXAMPLE_PROTO_1),
'test_1': ProtoUtil.message_to_any(self.EXAMPLE_PROTO_3),
})
self.assertEqual(proto_table_storage.get_num_entries(), 2)
gclient_ext.cp_file(self.TEST_DATA_2, self.TEST_DATA_4)
def _proto_table_storage_impl(self, request):
self._logger.info("Getting request of proto table storage read.")
read_params = dict(request.params)
if 'proto_module' in read_params:
read_params['message_type'] = ProtoUtil.infer_message_type_from_str(
message_type_str=read_params['message_type'],
modules=read_params['proto_module']
)
else:
read_params['message_type'] = ProtoUtil.infer_message_type_from_str(
message_type_str=read_params['message_type']
)
lru_key = (request.type, request.file_name)
storage = self._lru_cache_tool.get(key=lru_key)
if not storage:
self.sys_log("Did not find the storage in cache. Making a new one...")
storage = ProtoTableStorage()
storage.initialize_from_file(file_name=request.file_name)
self._lru_cache_tool.set(
key=lru_key,
value=storage
)
else:
self.sys_log("Found key in LRU cache.")
self._logger.info('Current cache size ' + str(self._lru_cache_tool.get_cur_capacity()))
read_params.pop('proto_module', None)
return storage.read(params=read_params)
def test_read_1(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
result_proto = proto_table_storage.read(params={
'key': 'test',
'message_type': NodeSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_1)
def test_write_1(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
proto_table_storage.write(data={'test': self.EXAMPLE_PROTO_1}, )
result_proto = proto_table_storage.read(params={
'key': 'test',
'message_type': NodeSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_1)
gclient_ext.cp_file(self.TEST_DATA_1, self.TEST_DATA_3)
def test_read_3(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
result_proto = proto_table_storage.read(params={'key': 'test1'})
self.assertIsNone(result_proto)
def test_read_2(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
result_proto = proto_table_storage.read(params={'key': 'test'})
self.assertEqual(result_proto,
ProtoUtil.message_to_any(self.EXAMPLE_PROTO_1))
def read_range(self, params):
def _reformat_time(timestamp):
if self.PARTITIONER_TYPE == PartitionerStorageType.YEARLY:
timestamp = timestamp.replace(month=1,
day=1,
hour=0,
minute=0,
second=0,
microsecond=0,
tzinfo=None)
elif self.PARTITIONER_TYPE == PartitionerStorageType.MONTHLY:
timestamp = timestamp.replace(day=1,
hour=0,
minute=0,
second=0,
microsecond=0,
tzinfo=None)
elif self.PARTITIONER_TYPE == PartitionerStorageType.DAILY:
timestamp = timestamp.replace(hour=0,
minute=0,
second=0,
microsecond=0,
tzinfo=None)
elif self.PARTITIONER_TYPE == PartitionerStorageType.HOURLY:
timestamp = timestamp.replace(minute=0,
second=0,
microsecond=0,
tzinfo=None)
else:
timestamp = timestamp.replace(second=0,
microsecond=0,
tzinfo=None)
return timestamp
assert 'start_time' in params and 'end_time' in params and params[
'start_time'] <= params['end_time']
while self._writer_status != Status.IDLE:
self.sys_log("Waiting for writer to finish.")
time.sleep(TimeSleepObj.ONE_SECOND)
self._reader_status = Status.RUNNING
oldest_dir, latest_dir = self.get_oldest_dir(), self.get_latest_dir()
if not latest_dir or not oldest_dir:
if self.is_empty():
self._logger.warning("Current partitioner [" +
self.get_dir_name() +
"] is empty, cannot read anything.")
self.sys_log("Current partitioner [" + self.get_dir_name() +
"] is empty, cannot read anything.")
return {}
oldest_dir = oldest_dir.replace(self._file_tree.get_root_name(), '')
latest_dir = latest_dir.replace(self._file_tree.get_root_name(), '')
oldest_timestamp = FileUtil.parse_dir_to_timestamp(dir_name=oldest_dir)
latest_timestamp = FileUtil.parse_dir_to_timestamp(dir_name=latest_dir)
start_time = max(_reformat_time(params['start_time']),
oldest_timestamp)
end_time = min(_reformat_time(params['end_time']), latest_timestamp)
result = {}
try:
while start_time <= end_time:
dir_list = FileUtil.parse_timestamp_to_dir(
timestamp=start_time).split('/')
dir_name = '/'.join(
dir_list[:self.PARTITIONER_TYPE_TO_HEIGHT_MAP[
self.PARTITIONER_TYPE]])
dir_name = FileUtil.join_paths_to_dir(
root_dir=self._file_tree.get_root_name(),
base_name=dir_name)
if FileUtil.does_dir_exist(dir_name=dir_name):
if self._underlying_storage.get_storage_type(
) == StorageType.PROTO_TABLE_STORAGE:
storage = ProtoTableStorage()
else:
storage = DefaultStorage()
file_names = FileUtil.list_files_in_dir(dir_name=dir_name)
for file_name in file_names:
storage.initialize_from_file(file_name=file_name)
if storage.get_storage_type(
) == StorageType.PROTO_TABLE_STORAGE:
result[file_name] = storage.read_all()
else:
result[file_name] = storage.read(
params={'num_line': -1})
if self.PARTITIONER_TYPE == PartitionerStorageType.YEARLY:
start_time = start_time.replace(year=start_time.year + 1,
month=1,
day=1)
elif self.PARTITIONER_TYPE == PartitionerStorageType.MONTHLY:
if start_time.month == 12:
start_time = start_time.replace(year=start_time.year +
1,
month=1,
day=1)
else:
start_time = start_time.replace(
month=start_time.month + 1)
elif self.PARTITIONER_TYPE == PartitionerStorageType.DAILY:
start_time += datetime.timedelta(days=1)
elif self.PARTITIONER_TYPE == PartitionerStorageType.HOURLY:
start_time += datetime.timedelta(hours=1)
else:
start_time += datetime.timedelta(minutes=1)
self._reader_status = Status.IDLE
return result
except Exception as err:
self.sys_log("Read range in dir [" + self.get_dir_name() +
"] got exception " + str(err) + '.')
self._logger.error("Read range in dir [" + self.get_dir_name() +
"] got exception " + str(err) + '.')
raise StorageReadException("Read range in dir [" +
self.get_dir_name() +
"] got exception " + str(err) + '.')