use window approach

sds_data_manager/lambda_code/SDSCode/api_lambdas/latest_version_query.py

@@ -0,0 +1,53 @@
+"""Science-file "latest version" resolution logic for the query API."""
+
+from collections.abc import Sequence
+
+from sqlalchemy import ColumnElement, Select, func, select
+
+from ..database.models import FILE_ID_COLUMNS, ScienceFiles
+
+
+def build_latest_version_query(
+    filters: Sequence[ColumnElement] = (),
+    major_only: bool = False,
+) -> Select:
+    """Build a query selecting the latest version of each science file.
+
+    Rows that share every :data:`FILE_ID_COLUMNS` value are just different
+    versions of the same file. The query uses a window function to rank them
+    by version and keeps only the ones with rank=1.
+
+    Parameters
+    ----------
+    filters : sequence of column expressions, optional
+        WHERE conditions applied *before* the version selection.
+    major_only : bool, optional
+        When True, return every minor version of each series' latest major
+        version instead of just the single newest file.
+
+    Returns
+    -------
+    Select
+        A SELECT of the :class:`ScienceFiles` table columns, restricted to the
+        latest version of each series.
+
+    """
+    table = ScienceFiles.__table__
+
+    partition_by = [table.c[column] for column in FILE_ID_COLUMNS]
+    order_by = [table.c.major_version.desc()]
+    if not major_only:
+        order_by.append(table.c.minor_version.desc())
+    rank = (
+        func.rank()
+        .over(partition_by=partition_by, order_by=order_by)
+        .label("version_rank")
+    )
+    ranked = select(table, rank).where(*filters).subquery()
+
+    rank_col = ranked.c[rank.name]
+    top_only = select(ranked).where(rank_col == 1)
+
+    # excludes the added RANK column
+    original_columns = [ranked.c[col.name] for col in table.c]
+    return top_only.with_only_columns(*original_columns)

sds_data_manager/lambda_code/SDSCode/api_lambdas/query_api.py

@@ -3,32 +3,47 @@
 import datetime
 import json
 import logging
-from collections import namedtuple
+from enum import StrEnum
 
-from sqlalchemy import and_, func, or_, select
+from sqlalchemy import func, select
 
 from ..api_lambdas.utils import is_authenticated_user
+from ..api_lambdas.latest_version_query import build_latest_version_query
 from ..database import database as db
 from ..database import models
 
 # Logger setup
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
-# Columns that, together with repointing, identify a unique science file
-# "series" when resolving the latest version.
-_VERSION_GROUPING_COLUMNS = (
-    "instrument",
-    "data_level",
-    "descriptor",
-    "start_date",
-)
+# Maps the `table` query param to its model.
+_TABLE_MODELS = {
+    "science": models.ScienceFiles,
+    "ancillary": models.AncillaryFiles,
+    "spice": models.SPICEFiles,
+    "quicklook": models.QuicklookFiles,
+}
+
+# Valid query parameters per table: its columns (minus "id"), plus "end_date"
+# for every table but "ancillary", plus the ingestion date range params.
+_VALID_PARAMETERS = {
+    table: [
+        *(column.key for column in model.__table__.columns if column.key != "id"),
+        *(["end_date"] if table != "ancillary" else []),
+        "ingestion_start_date",
+        "ingestion_end_date",
+    ]
+    for table, model in _TABLE_MODELS.items()
+}
+
 
-# The two ways a science query resolves "latest". NEWEST keeps only the single
-# newest file per series (latest major, then latest minor); LATEST_MAJOR keeps
-# every minor version of the latest major.
-LatestVersionMode = namedtuple("LatestVersionMode", ["newest", "latest_major"])
-LATEST_VERSION_MODE = LatestVersionMode(newest="newest", latest_major="latest_major")
+class LatestVersionMode(StrEnum):
+    """The two ways a science query resolves "latest"."""
+
+    # single newest file per series (latest major, then latest minor)
+    NEWEST = "newest"
+    # every minor version of the latest major
+    LATEST_MAJOR = "latest_major"
 
 
 def _parse_version_alias(value):
@@ -71,8 +86,8 @@ def _resolve_science_version_mode(query_params):
 
     Returns
     -------
-    str or None
-        A ``LATEST_VERSION_MODE`` value, or None when a concrete major_version
+    LatestVersionMode or None
+        A ``LatestVersionMode`` value, or None when a concrete major_version
         was requested (no latest restriction applied).
 
     Raises
@@ -98,87 +113,45 @@ def _resolve_science_version_mode(query_params):
     if "major_version" in query_params:
         return None
     if latest_flag:
-        return LATEST_VERSION_MODE.newest
-    return LATEST_VERSION_MODE.latest_major
-
+        return LatestVersionMode.NEWEST
+    return LatestVersionMode.LATEST_MAJOR
 
-def _latest_version_filters(model, mode, released_only):
-    """Build WHERE clauses restricting science results to the latest version.
 
-    "Latest" is resolved with correlated subqueries: for each candidate result
-    row, a subquery computes the maximum version within that row's file
-    "series" (same instrument / data_level / descriptor / start_date /
-    repointing) and the row is kept only if it matches that maximum.
+def _filter_condition(cols, param, value):
+    """Build the SQLAlchemy filter expression for one query parameter.
 
     Parameters
     ----------
-    model
-        The science table model.
-    mode : str
-        One of ``LATEST_VERSION_MODE``. ``latest_major`` keeps the latest major
-        version with ALL of its minor versions; ``newest`` keeps only the single
-        newest file (latest major and, within it, latest minor).
-    released_only : bool
-        When True, only released files count toward "latest". This is applied
-        *inside* the max subqueries (not as a plain filter on the outer query)
-        so that an unreleased newer version cannot hide the latest released
-        version from unauthenticated users.
+    cols
+        The column collection to filter against.
+    param : str
+        The query parameter name.
+    value : str
+        The query parameter value.
 
     Returns
     -------
-    list
-        SQLAlchemy boolean clauses to AND into the query.
+    ColumnElement
+        A SQLAlchemy boolean clause to AND into the query.
 
     """
-    outer = model.__table__
-
-    def same_series_as_outer(inner):
-        """Conditions correlating an inner alias to the outer row's series."""
-        # Equality on every grouping column that defines one file series.
-        conditions = [
-            getattr(inner.c, column) == getattr(outer.c, column)
-            for column in _VERSION_GROUPING_COLUMNS
-        ]
-        # repointing is nullable, so NULL-safe equality keeps NULL-repointing
-        # rows grouped together instead of dropping them from the correlation.
-        conditions.append(
-            or_(
-                inner.c.repointing == outer.c.repointing,
-                and_(inner.c.repointing.is_(None), outer.c.repointing.is_(None)),
+    match param:
+        case "start_date":
+            return cols.start_date >= datetime.datetime.strptime(value, "%Y%m%d")
+        case "end_date":
+            # TODO: Need to discuss as a team how to handle date queries. For now,
+            # the date queries will only look at the file start_date.
+            return cols.start_date <= datetime.datetime.strptime(value, "%Y%m%d")
+        case "ingestion_start_date":
+            return func.date(cols.ingestion_date) >= (
+                datetime.datetime.strptime(value, "%Y%m%d").date()
             )
-        )
-        # Only released files participate in the max when released_only is set.
-        if released_only:
-            conditions.append(inner.c.released.is_(True))
-        return conditions
-
-    # Keep rows whose major_version is the max major within their series.
-    # Each max subquery re-scans the science table, so it needs its own aliased
-    # reference to that table (a self-join): the alias is the inner scan, while
-    # the un-aliased `outer` is the row being tested, which the subquery
-    # correlates back to via same_series_as_outer().
-    major_inner = outer.alias("latest_major_inner")
-    max_major = (
-        select(func.max(major_inner.c.major_version))
-        .where(*same_series_as_outer(major_inner))
-        .scalar_subquery()
-    )
-    filters = [outer.c.major_version == max_major]
-
-    # For "newest", additionally keep only the max minor within that major.
-    if mode == LATEST_VERSION_MODE.newest:
-        # A second, independent self-join alias for the minor-version scan.
-        minor_inner = outer.alias("latest_minor_inner")
-        max_minor = (
-            select(func.max(minor_inner.c.minor_version))
-            .where(
-                *same_series_as_outer(minor_inner),
-                minor_inner.c.major_version == outer.c.major_version,
+        case "ingestion_end_date":
+            return func.date(cols.ingestion_date) <= (
+                datetime.datetime.strptime(value, "%Y%m%d").date()
             )
-            .scalar_subquery()
-        )
-        filters.append(outer.c.minor_version == max_minor)
-    return filters
+        case _:
+            return cols[param] == value
 
 
 def _format_search_results(search_results):
@@ -211,7 +184,7 @@ def _format_search_results(search_results):
     return search_results
 
 
-def lambda_handler(event, context):  # noqa: PLR0912
+def lambda_handler(event, context):
     """Entry point to the query API lambda.
 
     Parameters
@@ -230,35 +203,14 @@ def lambda_handler(event, context):  # noqa: PLR0912
 
     logger.info("Received event: " + json.dumps(event, indent=2))
 
-    TableModels = namedtuple(
-        "TableModels", ["science", "ancillary", "spice", "quicklook"]
-    )
-
-    table_models = TableModels(
-        science=models.ScienceFiles,
-        ancillary=models.AncillaryFiles,
-        spice=models.SPICEFiles,
-        quicklook=models.QuicklookFiles,
-    )
-
-    # add session, pick model like in indexer and add query to filter_as
     # Make a mutable copy so we can pre-process science version parameters.
     query_params = dict(event["queryStringParameters"])
-    # get desired table for query
-    query_table = query_params.get("table", "science")
-
+    query_table = query_params.pop("table", "science")
     logger.info(f"Querying table: {query_table}")
-    model = getattr(table_models, query_table)
 
-    # select the given table for the query
-    authenticated = is_authenticated_user(event)
-    query = select(model.__table__)
-    if not authenticated:
-        query = query.filter(model.released)
+    model = _TABLE_MODELS[query_table]
+    table_columns = model.__table__.c
 
-    # Science-only version handling: a backwards-compatible `version` alias plus
-    # server-side resolution of "latest". Other tables keep `version` as a real
-    # column and are left untouched.
     version_mode = None
     if query_table == "science":
         try:
@@ -271,26 +223,9 @@ def lambda_handler(event, context):  # noqa: PLR0912
                 ),
             }
 
-    # get a list of all valid search parameters
-    valid_parameters = [
-        column.key for column in model.__table__.columns if column.key not in ["id"]
-    ]
-    # Up until this point, valid_parameters are the same as the
-    # columns in the selected table. And looks like we removed
-    # the "id" column from the list. But we also need to add
-    # 'end_date' to the list of valid_parameters but only for
-    # the science table.
-    if query_table != "ancillary":
-        valid_parameters.append("end_date")
-    valid_parameters.append("ingestion_start_date")
-    valid_parameters.append("ingestion_end_date")
-
-    # go through each query parameter to set up sqlalchemy query conditions
+    valid_parameters = _VALID_PARAMETERS[query_table]
+    filters = []
     for param, value in query_params.items():
-        # skip the table parameter
-        if param == "table":
-            continue
-        # confirm that the query parameter is valid
         if param not in valid_parameters:
             response = {
                 "statusCode": 400,
@@ -304,45 +239,31 @@ def lambda_handler(event, context):  # noqa: PLR0912
                 "{query_table}, valid options are: {valid_parameters}"
             )
             return response
-        # check if we're search for start_date or end date or ingestion dates to
-        # setup the correct "where" time condition
-        if param == "start_date":
-            query = query.where(
-                model.start_date >= datetime.datetime.strptime(value, "%Y%m%d")
-            )
-        elif param == "end_date":
-            # TODO: Need to discuss as a team how to handle date queries. For now,
-            # the date queries will only look at the file start_date.
-            query = query.where(
-                model.start_date <= datetime.datetime.strptime(value, "%Y%m%d")
-            )
-        elif param == "ingestion_start_date":
-            # filtering by ingestion date
-            query = query.where(
-                func.date(model.ingestion_date)
-                >= datetime.datetime.strptime(value, "%Y%m%d").date()
-            )
-        elif param == "ingestion_end_date":
-            query = query.where(
-                func.date(model.ingestion_date)
-                <= datetime.datetime.strptime(value, "%Y%m%d").date()
-            )
-        # all non-time string matching parameters
-        else:
-            query = query.where(getattr(model, param) == value)
+        filters.append(_filter_condition(table_columns, param, value))
 
-    # Restrict science results to the latest version when no concrete
-    # major_version was requested. NEWEST keeps a single file per series;
-    # LATEST_MAJOR keeps all minor versions of the latest major.
-    if version_mode is not None:
-        for clause in _latest_version_filters(model, version_mode, not authenticated):
-            query = query.where(clause)
+    # if not authenticated, restrict to released only
+    authenticated = is_authenticated_user(event)
+    if not authenticated:
+        filters.append(table_columns.released)
+
+    if version_mode is None:
+        # if not filtering by version, simply SELECT ... WHERE ...
+        query = select(model.__table__).where(*filters)
+        cols = table_columns
+    else:
+        # otherwise, also include the rank subquery
+        query = build_latest_version_query(
+            filters=filters,
+            major_only=version_mode == LatestVersionMode.LATEST_MAJOR,
+        )
+        # important not to use table_columns hereafter
+        cols = query.selected_columns
 
     # We want to order the query returns by the filename
     # This will implicitly sort by: instrument, data level, descriptor, start_date, ...
     # Default for the table is by the ascending id so by insertion order
     # This fails for the SPICE table because it uses 'file_name'
-    query = query.order_by(model.file_path)
+    query = query.order_by(cols.file_path)
 
     with db.Session() as session:
         search_results = session.execute(query).all()