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 content_serializer(self, content):
op_content = OperatorContentPlain()
try:
op_content.plan_content.bytes_val = str(content)
except Exception as err:
self._logger.warning("Seralize content with error: " + str(err) + '.')
return ProtoUtil.message_to_any(message=op_content)
def write(self, data, params=None):
if not params:
params = {}
if 'overwrite' not in params:
params['overwrite'] = True
assert isinstance(data, dict)
try:
for key, val in data.items():
if not params['overwrite'] and key in self._table_message.data:
continue
any_message = ProtoUtil.message_to_any(message=val)
self._table_message.data[key].CopyFrom(any_message)
if len(self._table_message.data) > 1000:
self._SYS_LOGGER.warning(
"Warning: the table content is too large, considering using Partitioner "
"combined with proto table.")
self._table_message.updated_time = str(
TimezoneUtil.cur_time_in_pst())
self.increment_rpc_count_by(n=1)
FileUtil.write_proto_to_file(proto=self._table_message,
file_name=self._file_name)
except Exception as err:
self._SYS_LOGGER.error("Write to file [" + self.get_file_name() +
"] got exception: " + str(err) + '.')
self._logger.error("Write to file [" + self.get_file_name() +
"] got exception: " + str(err) + '.')
raise StorageWriteException("Write to file [" +
self.get_file_name() +
"] got exception: " + str(err) + '.')
def test_message_to_any(self):
node_snapshot = NodeSnapshot()
node_snapshot.node_name = 'test'
node_snapshot.children_names.append('children1')
node_snapshot.children_names.append('children2')
node_snapshot.parents_names.append('parent1')
node_snapshot.parents_names.append('parent2')
any_message = ProtoUtil.message_to_any(message=node_snapshot)
self.assertEqual(
ProtoUtil.any_to_message(message_type=NodeSnapshot,
any_message=any_message), node_snapshot)
def send_request(self, request):
self._response = None
generic_request = GenericRPCRequest()
generic_request.request_data.CopyFrom(
ProtoUtil.message_to_any(message=request))
generic_request.timestamp = str(TimezoneUtil.cur_time_in_pst())
self._corr_id = generic_request.uuid = str(uuid.uuid4())
if self.RESPONSE_MESSAGE_TYPE:
generic_request.message_type = ProtoUtil.infer_str_from_message_type(
message_type=self.RESPONSE_MESSAGE_TYPE)
self._SYS_LOGGER.info("Getting request of uuid [" +
generic_request.uuid + '] in queue [' +
self.get_queue_name() + '].')
self._logger.info("Getting request of uuid [" + generic_request.uuid +
'] in queue [' + self.get_queue_name() + '].')
try:
generic_request_str = ProtoUtil.message_to_string(
proto_message=generic_request)
self._channel.basic_publish(
exchange='',
routing_key=self._queue_name,
properties=pika.BasicProperties(reply_to=self._callback_queue,
correlation_id=self._corr_id),
body=base64.b64encode(generic_request_str))
wait_start_time = TimezoneUtil.cur_time_in_pst()
while not self._response:
self._connection.process_data_events(time_limit=int(
EnvUtil.get_pslx_env_variable('PSLX_QUEUE_TIMEOUT')))
if TimezoneUtil.cur_time_in_pst() - wait_start_time > \
datetime.timedelta(seconds=int(EnvUtil.get_pslx_env_variable('PSLX_QUEUE_TIMEOUT'))):
break
if not self.RESPONSE_MESSAGE_TYPE or self._response is None:
return None
else:
return ProtoUtil.any_to_message(
message_type=self.RESPONSE_MESSAGE_TYPE,
any_message=self._response.response_data)
except Exception as err:
self._logger.error('Queue [' + self.get_queue_name() +
"] send request with error " + str(err) + '.',
publish=True)
self._SYS_LOGGER.error('Queue [' + self.get_queue_name() +
"] send request with error " + str(err) +
'.')
def get_operator_snapshot(self, output_file=None):
snapshot = OperatorSnapshot()
snapshot.operator_name = self.get_node_name()
snapshot.data_model = self.get_data_model()
snapshot.status = self.get_status()
snapshot.node_snapshot.CopyFrom(self.get_node_snapshot())
snapshot.class_name = self.get_full_class_name()
if self._start_time:
snapshot.start_time = str(self._start_time)
if self.get_status() == Status.SUCCEEDED and self._end_time:
snapshot.end_time = str(self._end_time)
if self._persistent:
assert self.CONTENT_MESSAGE_TYPE is not None
snapshot.content.CopyFrom(ProtoUtil.message_to_any(message=self._content))
if output_file and self._config['save_snapshot'] and 'Dummy' not in self.get_class_name():
self.sys_log("Saved to file " + output_file + '.')
with FileLockTool(output_file, read_mode=False):
FileUtil.write_proto_to_file(
proto=snapshot,
file_name=output_file
)
return snapshot
def write(self, data, params=None):
if not params:
params = {}
if 'overwrite' not in params:
params['overwrite'] = True
while self._reader_status != Status.IDLE:
self.sys_log("Waiting for reader to finish.")
time.sleep(TimeSleepObj.ONE_SECOND)
while self._deleter_status != Status.IDLE:
self.sys_log("Waiting for deleter to finish.")
time.sleep(TimeSleepObj.ONE_SECOND)
self._writer_status = Status.RUNNING
assert isinstance(data, dict)
try:
for key, val in data.items():
if not params['overwrite'] and key in self._table_message.data:
continue
any_message = ProtoUtil.message_to_any(message=val)
self._table_message.data[key].CopyFrom(any_message)
if len(self._table_message.data) > 1000:
self.sys_log("Warning: the table content is too large, considering using Partitioner "
"combined with proto table.")
self._table_message.updated_time = str(TimezoneUtil.cur_time_in_pst())
with FileLockTool(self._file_name, read_mode=False):
FileUtil.write_proto_to_file(
proto=self._table_message,
file_name=self._file_name
)
self._writer_status = Status.IDLE
except Exception as err:
self.sys_log("Write to file [" + self.get_file_name() + "] got exception: " + str(err) + '.')
self._logger.error("Write to file [" + self.get_file_name() + "] got exception: " + str(err) + '.')
raise StorageWriteException("Write to file [" + self.get_file_name() + "] got exception: " + str(err) + '.')
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))