use window approach (#1502)

* use window approach

* fix _VALID_PARAMETERS setup

* revert ancillary file handling

* handle invalid table with error code 400

* handle invalid table with error code 400

* style: apply ruff docstring and format fixes

Assisted-by: Claude

* revert release_api changes; upstream #1533 supersedes them

Assisted-by: Claude

Author: hafarooki

sds_data_manager/lambda_code/SDSCode/api_lambdas/query_api.py

@@ -3,32 +3,53 @@
 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.utils import build_latest_version_query, is_authenticated_user
 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 include...
+#   all table columns
+# + ingestion_start_date/ingestion_end_date,
+# + "end_date" for tables with a start_date but no end_date
+#   (science/quicklook have start_date only; ancillary has both, spice has neither).
+_VALID_PARAMETERS = {
+    table: [
+        *model.__table__.c.keys(),
+        *(
+            ["end_date"]
+            if "start_date" in model.__table__.c and "end_date" not in model.__table__.c
+            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 +92,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 +119,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 +190,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 +209,23 @@ 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)
+    if query_table not in _TABLE_MODELS:
+        return {
+            "statusCode": 400,
+            "body": json.dumps(
+                f"{query_table} is not a valid table. "
+                f"Valid tables are: {list(_TABLE_MODELS)}"
+            ),
+        }
+
+    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 +238,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 +254,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()

tests/conftest.py

@@ -1,10 +1,39 @@
 """Setup items for all test types."""
 
 import os
+from contextlib import contextmanager
+from unittest.mock import patch
 
 import boto3
 import pytest
 from moto import mock_dynamodb
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+from sds_data_manager.lambda_code.SDSCode.database import database as db
+from sds_data_manager.lambda_code.SDSCode.database.models import Base
+
+
+@contextmanager
+def in_memory_session(engine=None):
+    """Yield an in-memory SQLite session with ``db.Session`` patched to it.
+
+    Shared scaffolding for the DB-backed ``session`` fixtures. Pass a pre-built
+    ``engine`` (e.g. with event listeners already attached) or let it build a
+    default in-memory engine.
+    """
+    if engine is None:
+        engine = create_engine("sqlite:///:memory:")
+    Base.metadata.create_all(engine)
+    with patch.object(db, "Session") as mock_session:
+        session = sessionmaker(bind=engine)()
+        mock_session.return_value = session
+        try:
+            yield session
+        finally:
+            session.rollback()
+            session.close()
+            Base.metadata.drop_all(engine)
 
 
 @pytest.fixture