sdss_id list from search string

python/valis/db/queries.py

@@ -6,6 +6,7 @@
 
 from typing import Union
 import peewee
+from peewee import fn
 import astropy.units as u
 from astropy.coordinates import SkyCoord
 from sdssdb.peewee.sdss5db import vizdb
@@ -119,6 +120,100 @@ def cone_search(ra: Union[str, float], dec: Union[str, float],
                                               dec_col='dec_sdss_id'))
 
 
+def get_paged_sdss_id_list(search_integer: int, page_number: int = 1, items_per_page: int = 10) -> peewee.ModelSelect:
+    """ Return a paged list of sdss_id values.
+
+    Return paginated and ordered sdss_id column values from the 
+    vizdb.SDSSidStacked table using the peewee ORM.
+    We return the peewee ModelSelect
+    directly here so it can be easily combined with other queries, 
+    if needed. 
+
+    Parameters
+    ----------
+    search_string : str
+        String that matches the starting digits of the returned sdss_id values. If empty, no matching is applied.
+    page_number : int
+        Page number of the returned sdss_id values.
+    items_per_page : int
+        Number of sdss_id values displayed in the page.
+
+    Returns
+    -------
+    peewee.ModelSelect
+        the ORM query
+    """
+
+    max_sdss_id = vizdb.SDSSidStacked.select(fn.MAX(vizdb.SDSSidStacked.sdss_id)).scalar()
+
+    max_num_digits = len(str(max_sdss_id))
+    num_search_digits = len(str(search_integer))
+    max_i =  max_num_digits - num_search_digits + 1
+
+    where_condition = (False)
+    for i in range(0, max_i):
+        min_id = int(search_integer * 10**(i))
+        max_id = int((search_integer + 1) * 10**(i))
+        where = ((vizdb.SDSSidStacked.sdss_id >=  min_id) & (vizdb.SDSSidStacked.sdss_id < max_id))
+        where_condition = where_condition | where
+
+    return vizdb.SDSSidStacked.select(vizdb.SDSSidStacked.sdss_id)\
+                              .where(where_condition)\
+                              .order_by(vizdb.SDSSidStacked.sdss_id)\
+                              .paginate(page_number, items_per_page)
+
+
+def get_targets_by_sdss_id(sdss_id: int) -> peewee.ModelSelect:
+    """ Perform a search for SDSS targets on vizdb.SDSSidStacked based on the sdss_id.
+
+    Perform a search for SDSS targets using the peewee ORM in the
+    vizdb.SDSSidStacked table. We return the peewee ModelSelect
+    directly here so it can be easily combined with other queries, 
+    if needed.
+
+    In the route endpoint itself, remember to return wrap this in a list.
+
+    Parameters
+    ----------
+    sdss_id : int
+        the sdss_id
+
+    Returns
+    -------
+    peewee.ModelSelect
+        the ORM query
+    """
+
+    return vizdb.SDSSidStacked.select().where(vizdb.SDSSidStacked.sdss_id == sdss_id)
+
+
+def get_targets_by_catalog_id(catalog_id: int) -> peewee.ModelSelect:
+    """ Perform a search for SDSS targets on vizdb.SDSSidStacked based on the catalog_id.
+
+    Perform a search for SDSS targets using the peewee ORM in the
+    vizdb.SDSSidStacked table. We return the peewee ModelSelect
+    directly here so it can be easily combined with other queries, 
+    if needed.
+
+    In the route endpoint itself, remember to return wrap this in a list.
+
+    Parameters
+    ----------
+    catalog_id : int
+        the catalog_id
+
+    Returns
+    -------
+    peewee.ModelSelect
+        the ORM query
+    """
+
+    return vizdb.SDSSidStacked.select()\
+                              .join(vizdb.SDSSidFlat, on=(vizdb.SDSSidStacked.sdss_id == 
+                                                          vizdb.SDSSidFlat.sdss_id))\
+                              .where(vizdb.SDSSidFlat.catalogid == catalog_id)
+
+
 def carton_program_list(name_type: str) -> peewee.ModelSelect:
     """
     Return a list of either all cartons or programs from targetdb
@@ -184,4 +279,4 @@ def carton_program_search(name: str, name_type: str) -> peewee.ModelSelect:
                             .where(getattr(targetdb.Carton, name_type) == name)
     model_stack = vizdb.SDSSidStacked.select()\
                                      .join(model, on=(model.c.sdss_id == vizdb.SDSSidStacked.sdss_id))
-    return model_stack
+    return model_stack

python/valis/routes/query.py

@@ -12,7 +12,9 @@
 from valis.db.db import get_pw_db
 from valis.db.models import SDSSidStackedBase, SDSSidPipesBase
 from valis.db.queries import (cone_search, append_pipes, carton_program_search,
-                              carton_program_list, carton_program_map)
+                              carton_program_list, carton_program_map, 
+                              get_targets_by_sdss_id, get_targets_by_catalog_id,
+                              get_paged_sdss_id_list)
 
 
 class SearchCoordUnits(str, Enum):
@@ -86,6 +88,29 @@ async def cone_search(self,
         """ Perform a cone search """
         return list(cone_search(ra, dec, radius, units=units))
 
+    @router.get('/sdssid', summary='Perform a search for an SDSS target based on the sdss_id',
+                response_model=Union[SDSSidStackedBase, dict], dependencies=[Depends(get_pw_db)])
+    async def sdss_id_search(self, sdss_id: Union[int, str] = Query(..., description='Value of sdss_id', example=47510284)):
+        """ Perform an sdss_id search. Assumes a maximum of one target per sdss_id. Empty object returned when no match is found."""
+        targets = get_targets_by_sdss_id(int(sdss_id))
+        return targets[0] if len(targets) > 0 else {}
+
+    @router.get('/catalogid', summary='Perform a search for SDSS targets based on the catalog_id',
+                response_model=List[SDSSidStackedBase], dependencies=[Depends(get_pw_db)])
+    async def catalog_id_search(self, catalog_id: Union[int, str] = Query(..., description='Value of catalog_id', example=7613823349)):
+        """ Perform a catalog_id search """
+        return list(get_targets_by_catalog_id(int(catalog_id)))
+
+    @router.get('/list/sdssid', summary='Return a paged list of sdss_id values',
+                response_model=List[int], dependencies=[Depends(get_pw_db)])
+    async def sdss_id_list(self, search_integer: int = Query(default="", description='Integer that matches the starting digits of the returned sdss_id values', ge=1, example=1),
+                                 page_number: int = Query(..., description='Page number of the returned items', gt=0, example=1),
+                                 items_per_page: int = Query(..., description='Number of items displayed in a page', gt=0, example=10)):
+        """ Return an ordered and paged list of sdss_id values."""
+        sdss_ids = get_paged_sdss_id_list(search_integer, page_number, items_per_page)
+        return list(sdss_ids.scalars())
+
+
     @router.get('/list/cartons', summary='Return a list of all cartons',
                 response_model=list, dependencies=[Depends(get_pw_db)])
     async def cartons(self,
@@ -116,3 +141,4 @@ async def carton_program(self,
                                                     description='Specify search on carton or program', example='carton')):
         """ Perform a search on carton or program """
         return list(carton_program_search(name, name_type))
+