diff --git a/python/valis/db/queries.py b/python/valis/db/queries.py index 000a8e5..0590682 100644 --- a/python/valis/db/queries.py +++ b/python/valis/db/queries.py @@ -4,6 +4,7 @@ # all resuable queries go here +from contextlib import contextmanager import itertools import packaging import uuid @@ -61,16 +62,15 @@ def append_pipes(query: peewee.ModelSelect, table: str = 'stacked', # Run initial query as a temporary table. temp = create_temporary_table(query, indices=['sdss_id']) - model = vizdb.SDSSidStacked if table == 'stacked' else vizdb.SDSSidFlat qq = temp.select(temp.__star__, - vizdb.SDSSidToPipes.in_boss, - vizdb.SDSSidToPipes.in_apogee, - vizdb.SDSSidToPipes.in_bvs, - vizdb.SDSSidToPipes.in_astra, - vizdb.SDSSidToPipes.has_been_observed, - vizdb.SDSSidToPipes.release, - vizdb.SDSSidToPipes.obs, - vizdb.SDSSidToPipes.mjd).\ + vizdb.SDSSidToPipes.in_boss, + vizdb.SDSSidToPipes.in_apogee, + vizdb.SDSSidToPipes.in_bvs, + vizdb.SDSSidToPipes.in_astra, + vizdb.SDSSidToPipes.has_been_observed, + vizdb.SDSSidToPipes.release, + vizdb.SDSSidToPipes.obs, + vizdb.SDSSidToPipes.mjd).\ join(vizdb.SDSSidToPipes, on=(temp.c.sdss_id == vizdb.SDSSidToPipes.sdss_id)).\ distinct(temp.c.sdss_id) @@ -292,7 +292,14 @@ def carton_program_search(name: str, """ if query is None: - query = vizdb.SDSSidStacked.select(vizdb.SDSSidStacked) + query = vizdb.SDSSidStacked.select(vizdb.SDSSidStacked).distinct() + + # NOTE: These setting seem to help when querying some cartons or programs, mainly + # those with small number of targets, and in some cases with these the query + # actually applies the LIMIT more efficiently, but it's not a perfect solution. + vizdb.database.execute_sql('SET enable_gathermerge = off;') + vizdb.database.execute_sql('SET parallel_tuple_cost = 100;') + vizdb.database.execute_sql('SET enable_bitmapscan = off;') query = (query.join( vizdb.SDSSidFlat, @@ -944,7 +951,8 @@ def get_target_by_altid(id: str | int, idtype: str = None) -> peewee.ModelSelect return get_targets_by_sdss_id(res.sdss_id) -def create_temporary_table(query: peewee.ModelSelect, indices: list[str] | None = None) -> peewee.Table: +def create_temporary_table(query: peewee.ModelSelect, + indices: list[str] | None = None) -> Generator[None, None, peewee.Table]: """Create a temporary table from a query.""" table_name = uuid.uuid4().hex[0:8] diff --git a/python/valis/routes/query.py b/python/valis/routes/query.py index 9c2e24a..866430c 100644 --- a/python/valis/routes/query.py +++ b/python/valis/routes/query.py @@ -41,6 +41,7 @@ class SearchModel(BaseModel): program: Optional[str] = Field(None, description='The program name', example='bhm_rm') carton: Optional[str] = Field(None, description='The carton name', example='bhm_rm_core') observed: Optional[bool] = Field(True, description='Flag to only include targets that have been observed', example=True) + limit: Optional[int] = Field(None, description='Limit the number of returned targets', example=100) class MainResponse(SDSSModel): """ Combined model from all individual query models """ @@ -105,6 +106,13 @@ async def main_search(self, body: SearchModel): query = carton_program_search(body.program or body.carton, 'program' if body.program else 'carton', query=query) + + # DANGER!!! This limit applies *before* the append_pipes call. If the + # append_pipes call includes observed=True we may have limited things in + # such a way that only unobserved or very few targets are returned. + if body.limit: + query = query.limit(body.limit) + # append query to pipes if query: query = append_pipes(query, observed=body.observed)