Skip to content

use window approach #1502

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tmplummer merged 8 commits into from
Jul 30, 2026
Merged

use window approach #1502

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0487ea3
use window approach
hafarooki Jul 18, 2026
16773a1
fix _VALID_PARAMETERS setup
hafarooki Jul 20, 2026
886dd15
revert ancillary file handling
hafarooki Jul 21, 2026
b160dcd
handle invalid table with error code 400
hafarooki Jul 21, 2026
7dae958
handle invalid table with error code 400
hafarooki Jul 21, 2026
6dd0b4e
Merge upstream/dev into use-window-instead-of-subquery
hafarooki Jul 30, 2026
2c2212c
style: apply ruff docstring and format fixes
hafarooki Jul 30, 2026
b2032cd
revert release_api changes; upstream #1533 supersedes them
hafarooki Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
264 changes: 100 additions & 164 deletions sds_data_manager/lambda_code/SDSCode/api_lambdas/query_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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,

@tmplummer tmplummer Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this specific to ancillary files that can have a start and end date? I would think that the desired results would be to provide any ancillary files that would be used for the full range that is queried.

@hafarooki hafarooki Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure tbh. That TODO comment was in the code from before my changes

@hafarooki hafarooki Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can define our own desired behavior here since it was not defined previously

# 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):
Expand Down Expand Up @@ -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
Expand All @@ -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
Comment thread
hafarooki marked this conversation as resolved.

# 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:
Expand All @@ -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,
Expand All @@ -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)
Comment on lines 277 to +281

@hafarooki hafarooki Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The problematic comment is from May 2025. It doesn't appear to be true of the latest codebase. I think I need input here on what the expected behavior should be.


with db.Session() as session:
search_results = session.execute(query).all()
Expand Down
29 changes: 29 additions & 0 deletions tests/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Loading
Loading