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))