From d13ae039c6b84735c96d2af4c009040a13a00873 Mon Sep 17 00:00:00 2001
From: shoulinwei <weishoulin@astrolab.cn>
Date: Thu, 5 Jan 2023 15:12:04 +0800
Subject: [PATCH] add columns

---
 csst_dfs_api/common/catalog.py      | 10 +++++++---
 csst_dfs_api/common/utils.py        | 31 ++++++++++++++++++++++++++++-
 tests/common/test_common_catalog.py | 10 +++++-----
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/csst_dfs_api/common/catalog.py b/csst_dfs_api/common/catalog.py
index 4398962..39e2943 100644
--- a/csst_dfs_api/common/catalog.py
+++ b/csst_dfs_api/common/catalog.py
@@ -9,13 +9,14 @@ class CatalogApi(object):
         self.module = Delegate().load(sub_module = "common")
         self.stub = getattr(self.module, "CatalogApi")()
     
-    def catalog_query(self, ra: float, dec: float, radius: float, catalog_name: str, min_mag: float,  max_mag: float,  obstime: int, limit: int):
+    def catalog_query(self, ra: float, dec: float, radius: float, catalog_name: str, columns: tuple, min_mag: float,  max_mag: float,  obstime: int, limit: int):
         ''' retrieval catalog
 
         :param ra: in deg
         :param dec:  in deg
         :param radius:  in deg
         :param catalog_name: one of ['gaia3','','']
+        :param columns: tuple of str, like ('ra','dec','phot_g_mean_mag')
         :param min_mag: minimal magnitude
         :param max_mag: maximal magnitude
         :param obstime: seconds  
@@ -32,12 +33,13 @@ class CatalogApi(object):
     def to_table(self, query_result):
         return to_fits_table(query_result)
 
-    def gaia3_query(self, ra: float, dec: float, radius: float, min_mag: float,  max_mag: float,  obstime: int, limit: int):
+    def gaia3_query(self, ra: float, dec: float, radius: float, columns: tuple,min_mag: float,  max_mag: float,  obstime: int, limit: int):
         """retrieval GAIA EDR 3
         
         :param ra: in deg
         :param dec:  in deg
         :param radius:  in deg
+        :param columns: tuple of str, like ('ra','dec','phot_g_mean_mag')
         :param min_mag: minimal magnitude
         :param max_mag: maximal magnitude
         :param obstime: seconds  
@@ -46,6 +48,8 @@ class CatalogApi(object):
         :returns: csst_dfs_common.models.Result
         """ 
         try:
-            return self.stub.gaia3_query(ra, dec, radius, min_mag, max_mag, obstime, limit)
+            if not columns:
+                raise Exception("columns is empty")
+            return self.stub.gaia3_query(ra, dec, radius, columns, min_mag, max_mag, obstime, limit)
         except Exception as e:
             return Result.error(message=repr(e))
diff --git a/csst_dfs_api/common/utils.py b/csst_dfs_api/common/utils.py
index 524f56c..4f8b2b6 100644
--- a/csst_dfs_api/common/utils.py
+++ b/csst_dfs_api/common/utils.py
@@ -80,7 +80,36 @@ def fields_dtypes(rec):
     dtypes = tuple(dtypes)
     return fields, dtypes
 
+def tuple_fields_dtypes(rec: tuple):
+    dtypes = []
+    for _, f in rec:
+        if f.type == int:
+            dtypes.append('i8')
+        elif f.type == float:
+            dtypes.append('f8')        
+        elif f.type == str:
+            dtypes.append('S2')
+        elif f.type == list:
+            dtypes.append('(12,)f8')       
+        else:
+            dtypes.append('S2')                
+    dtypes = tuple(dtypes)
+    return dtypes
+
 def to_table(query_result):
+    if not query_result.success or not query_result.data:
+        return Table()
+    fields = query_result['columns']
+    dtypes = tuple_fields_dtypes(query_result.data[0])
+    t = Table(names = fields, dtype = dtypes)
+    t.meta['comments'] = [str(query_result.data[0].__class__)]
+    t.meta['total'] = query_result['totalCount']
+
+    for rec in query_result.data:
+        t.add_row(rec)
+    return t
+
+def object_list_to_table(query_result):
     if not query_result.success or not query_result.data:
         return Table()
     fields, dtypes = fields_dtypes(query_result.data[0])
@@ -90,4 +119,4 @@ def to_table(query_result):
 
     for rec in query_result.data:
         t.add_row(tuple([rec.__getattribute__(k) for k in fields]))
-    return t
\ No newline at end of file
+    return t    
\ No newline at end of file
diff --git a/tests/common/test_common_catalog.py b/tests/common/test_common_catalog.py
index 2859250..b3331b5 100644
--- a/tests/common/test_common_catalog.py
+++ b/tests/common/test_common_catalog.py
@@ -13,18 +13,18 @@ class CommonCatalogTestCase(unittest.TestCase):
     def test_catalog_query(self):
         t= time.time()
         result = self.api.catalog_query(
-            ra=193.299, 
-            dec=19.08851, 
-            radius=2, 
+            ra=90, 
+            dec=24.5, 
+            radius=1, 
             catalog_name='gaia3', 
             min_mag=-1, 
             max_mag=-1, 
             obstime = -1, 
             limit = 0
         )
-        print(result)
+        # print(result)
         # dt = self.api.to_table(result)
         # df = dt.to_pandas()
         # print(df.head())
         print('used:', time.time()-t)
-        # print('return:', result['totalCount'\])
+        print('return:', len(result.data))
-- 
GitLab