diff --git a/opac_ssm_api/client.py b/opac_ssm_api/client.py index 186eedc..ed95f14 100644 --- a/opac_ssm_api/client.py +++ b/opac_ssm_api/client.py @@ -24,11 +24,11 @@ MAX_SEND_MESSAGE_LENGTH = int(os.getenv('MAX_SEND_MESSAGE_LENGTH', 90 * 1024 * 1024)) # 90MB try: - from opac_ssm_api import opac_pb2_grpc + from opac_ssm_api import opac_pb2_grpc, opac_pb2 except ImportError: logger.warning("Retrieving proto file from URL: http://%s:%s%s", HOST_PROTO_NAME, HTTP_PROTO_PORT, PROTO_PATH) utils.generate_pb_files(host=HOST_PROTO_NAME, port=HTTP_PROTO_PORT, proto_path=PROTO_PATH) - from opac_ssm_api import opac_pb2_grpc + from opac_ssm_api import opac_pb2_grpc, opac_pb2 class Client(object): @@ -129,7 +129,7 @@ def add_asset(self, pfile, filename='', filetype='', metadata='', logger.error(error_msg, pfile) raise IOError(error_msg) - asset = opac_pb2_grpc.Asset( + asset = opac_pb2.Asset( file=file_content, filename=filename, type=filetype, @@ -157,7 +157,7 @@ def get_asset(self, _id): logger.exception(msg) raise ValueError(msg) try: - asset = self.stubAsset.get_asset(opac_pb2_grpc.TaskId(id=_id)) + asset = self.stubAsset.get_asset(opac_pb2.TaskId(id=_id)) except Exception as e: logger.error(e) return (False, {'error_message': e.details()}) @@ -216,7 +216,7 @@ def query_asset(self, filters=None, metadata=None): else: raise ValueError("Metadada must be a dict or str") - assets = self.stubAsset.query(opac_pb2_grpc.Asset(**filters)).assets + assets = self.stubAsset.query(opac_pb2.Asset(**filters)).assets ret_list = [] @@ -258,7 +258,7 @@ def get_bucket(self, _id): logger.exception(msg) raise ValueError(msg) try: - bucket = self.stubAsset.get_bucket(opac_pb2_grpc.TaskId(id=_id)) + bucket = self.stubAsset.get_bucket(opac_pb2.TaskId(id=_id)) except Exception as e: logger.error(e) return (False, {'error_message': e.details()}) @@ -284,7 +284,7 @@ def get_asset_info(self, _id): raise ValueError(msg) try: - asset_info = self.stubAsset.get_asset_info(opac_pb2_grpc.TaskId(id=_id)) + asset_info = self.stubAsset.get_asset_info(opac_pb2.TaskId(id=_id)) except Exception as e: logger.error(e) return (False, {'error_message': e.details()}) @@ -309,7 +309,7 @@ def get_task_state(self, _id): logger.exception(msg) raise ValueError(msg) - task_state = self.stubAsset.get_task_state(opac_pb2_grpc.TaskId(id=_id)) + task_state = self.stubAsset.get_task_state(opac_pb2.TaskId(id=_id)) return task_state.state @@ -377,7 +377,7 @@ def update_asset(self, uuid, pfile=None, filename=None, filetype=None, metadata= if bucket_name: update_params['bucket'] = bucket_name - asset = opac_pb2_grpc.Asset(**update_params) + asset = opac_pb2.Asset(**update_params) return self.stubAsset.update_asset(asset).id else: @@ -397,8 +397,8 @@ def remove_asset(self, _id): if not isinstance(_id, six.string_types): raise ValueError('Param "_id" must be a str|unicode.') - if self.stubAsset.exists_asset(opac_pb2_grpc.TaskId(id=_id)): - return self.stubAsset.remove_asset(opac_pb2_grpc.TaskId(id=_id)) + if self.stubAsset.exists_asset(opac_pb2.TaskId(id=_id)): + return self.stubAsset.remove_asset(opac_pb2.TaskId(id=_id)) def add_bucket(self, name): """