import grpc import pickle import logging import io from csst_dfs_commons.models import Result from csst_dfs_proto.common.ephem import ephem_pb2, ephem_pb2_grpc from .service import get_grpc_channel from .utils import get_auth_headers log = logging.getLogger('csst') class CatalogApi(object): def __init__(self): self.stub = None def catalog_query(self, **kwargs): try: with get_grpc_channel() as c: self.stub = ephem_pb2_grpc.EphemSearchSrvStub(c) datas = io.BytesIO() totalCount = 0 resps = self.stub.Search(ephem_pb2.SearchRequest( conditions = { k:str(v) for k,v in kwargs.items() } ),metadata = get_auth_headers()) for resp in resps: if resp.success: datas.write(resp.records) totalCount = resp.totalCount else: return Result.error(message = str(resp.error.detail)) datas.flush() records = pickle.loads(datas.getvalue()) records, cols = records[0], records[1] columns = [] for col in cols: if col in columns: columns.append("%s_1" % (col, )) else: columns.append(col) return Result.ok_data(data = records).append("totalCount", totalCount).append("columns", columns) except grpc.RpcError as e: return Result.error(message="%s:%s" % (e.code().value, e.details()))