C9

csst_dfs_api_cluster/common/catalog.py

@@ -6,32 +6,32 @@ import io
 from csst_dfs_commons.models import Result
 
 from csst_dfs_proto.common.ephem import ephem_pb2, ephem_pb2_grpc
-from .service import grpc_channel
+from .service import get_grpc_channel
 from .utils import get_auth_headers
 
 log = logging.getLogger('csst')
 class CatalogApi(object):
     def __init__(self):
-        self.stub_class = ephem_pb2_grpc.EphemSearchSrvStub
         self.stub = None
 
-    @grpc_channel
     def catalog_query(self, **kwargs):
         try:
-            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())
+            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())
 
-            return Result.ok_data(data = records).append("totalCount", totalCount).append("columns", kwargs['columns'])
+                return Result.ok_data(data = records).append("totalCount", totalCount).append("columns", kwargs['columns'])
         except grpc.RpcError as e:
             return Result.error(message="%s:%s" % (e.code().value, e.details()))