catalog.py

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