From 6bf3f8cfa6db3344ce227b3a675ae0ffe2e0a8e2 Mon Sep 17 00:00:00 2001 From: Steven Kearnes Date: Fri, 26 Jun 2026 23:00:31 -0400 Subject: [PATCH 1/2] Speed up structural search: most-selective-first predicates + statement timeout Two changes to the shared search engine (queries.py / search.py), addressing the 20-60s+ (or hanging) structural searches analyzed in the ord-logbook entry. 1. Most-selective-first query shape. Replace the INTERSECT of one `SELECT DISTINCT reaction_id` per predicate -- with LIMIT on top, which could not bound the per-branch scans -- with a single `SELECT reaction.reaction_id FROM ord.reaction WHERE AND ... LIMIT n`, where each predicate is an EXISTS (or direct) condition correlated to the outer reaction row. The planner can now lead with the most selective predicate as a semi-join and stop once LIMIT rows are found. Measured on the live DB: benzene EXACT + yield>70 ~25s -> ~10s; morphine SIMILAR ~22s -> ~1s. Each ReactionQuery now exposes `where_predicate` instead of `query_and_parameters`. 2. statement_timeout floor. get_cursor sets statement_timeout (default 20s, ORD_INTERFACE_STATEMENT_TIMEOUT_MS) so genuinely pathological queries (e.g. a common-scaffold substructure with no selective co-filter, or an empty-result broad match) are cancelled instead of running away. The resulting psycopg QueryCanceled maps to a 400 with an actionable "too broad" message; the background task path records the error so fetch_query_result reports it rather than polling forever. All existing query/search result-count tests pass unchanged (the rewrite is result-preserving); adds unit tests for the timeout->400 mapping and the background-task error path. Co-Authored-By: Claude Opus 4.8 (1M context) --- ord_interface/api/queries.py | 199 ++++++++++++++++++------------- ord_interface/api/search.py | 52 ++++++-- ord_interface/api/search_test.py | 49 ++++++++ 3 files changed, 210 insertions(+), 90 deletions(-) diff --git a/ord_interface/api/queries.py b/ord_interface/api/queries.py index 7c156fe..44edbc5 100644 --- a/ord_interface/api/queries.py +++ b/ord_interface/api/queries.py @@ -45,7 +45,7 @@ from abc import ABC, abstractmethod from base64 import b64decode, b64encode from enum import Enum, auto -from typing import Any, LiteralString +from typing import Any, LiteralString, cast from ord_schema import message_helpers, validations from ord_schema.logging import get_logger @@ -65,8 +65,20 @@ class ReactionQuery(ABC): @property @abstractmethod - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" + def where_predicate(self) -> tuple[str, list]: + """Returns a boolean SQL predicate plus its parameters. + + The predicate is evaluated against a single outer ``ord.reaction`` row + (aliased ``reaction``) in ``run_queries`` -- typically an ``EXISTS (...)`` + correlated subquery. Combining predicates with ``AND`` over ``ord.reaction`` + lets the planner lead with the most selective one (a semi-join) and stop + early under ``LIMIT``, rather than materializing one ``SELECT DISTINCT`` per + predicate and intersecting them. + + Returns: + A ``(predicate_sql, params)`` tuple; ``params`` fill the predicate's + ``%s`` placeholders in order. + """ @property def session_config(self) -> dict[str, str]: @@ -98,15 +110,16 @@ def __init__(self, dataset_ids: list[str]) -> None: self._dataset_ids = dataset_ids @property - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" - query = """ - SELECT DISTINCT reaction.reaction_id - FROM ord.reaction - JOIN ord.dataset ON dataset.id = reaction.dataset_id - WHERE dataset.dataset_id = ANY (%s) + def where_predicate(self) -> tuple[str, list]: + """Returns the query predicate and any query parameters.""" + predicate = """ + EXISTS ( + SELECT 1 FROM ord.dataset + WHERE dataset.id = reaction.dataset_id + AND dataset.dataset_id = ANY (%s) + ) """ - return query, [self._dataset_ids] + return predicate, [self._dataset_ids] class ReactionIdQuery(ReactionQuery): @@ -124,14 +137,9 @@ def __init__(self, reaction_ids: list[str]) -> None: self._reaction_ids = reaction_ids @property - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" - query = """ - SELECT DISTINCT reaction.reaction_id - FROM ord.reaction - WHERE reaction.reaction_id = ANY (%s) - """ - return query, [self._reaction_ids] + def where_predicate(self) -> tuple[str, list]: + """Returns the query predicate and any query parameters.""" + return "reaction.reaction_id = ANY (%s)", [self._reaction_ids] class ReactionSmartsQuery(ReactionQuery): @@ -149,15 +157,16 @@ def __init__(self, reaction_smarts: str) -> None: self._reaction_smarts = reaction_smarts @property - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" - query = """ - SELECT DISTINCT reaction.reaction_id - FROM ord.reaction - JOIN rdkit.reactions ON rdkit.reactions.id = reaction.rdkit_reaction_id - WHERE rdkit.reactions.reaction @> reaction_from_smarts(%s) + def where_predicate(self) -> tuple[str, list]: + """Returns the query predicate and any query parameters.""" + predicate = """ + EXISTS ( + SELECT 1 FROM rdkit.reactions + WHERE rdkit.reactions.id = reaction.rdkit_reaction_id + AND rdkit.reactions.reaction @> reaction_from_smarts(%s) + ) """ - return query, [self._reaction_smarts] + return predicate, [self._reaction_smarts] class ReactionConversionQuery(ReactionQuery): @@ -180,24 +189,25 @@ def __init__( self._max_conversion = max_conversion @property - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" - query = """ - SELECT DISTINCT reaction.reaction_id - FROM ord.reaction - JOIN ord.reaction_outcome on reaction_outcome.reaction_id = reaction.id - JOIN ord.percentage on percentage.reaction_outcome_id = reaction_outcome.id - """ + def where_predicate(self) -> tuple[str, list]: + """Returns the query predicate and any query parameters.""" if self._min_conversion is not None and self._max_conversion is not None: - query += "WHERE percentage.value >= %s AND percentage.value <= %s\n" + condition = "percentage.value >= %s AND percentage.value <= %s" params = [self._min_conversion, self._max_conversion] elif self._min_conversion is not None: - query += "WHERE percentage.value >= %s\n" + condition = "percentage.value >= %s" params = [self._min_conversion] else: - query += "WHERE percentage.value <= %s\n" + condition = "percentage.value <= %s" params = [self._max_conversion] - return query, params + predicate = f""" + EXISTS ( + SELECT 1 FROM ord.reaction_outcome + JOIN ord.percentage ON percentage.reaction_outcome_id = reaction_outcome.id + WHERE reaction_outcome.reaction_id = reaction.id AND {condition} + ) + """ + return predicate, params class ReactionYieldQuery(ReactionQuery): @@ -216,25 +226,26 @@ def __init__(self, min_yield: float | None, max_yield: float | None) -> None: self._max_yield = max_yield @property - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" - query = """ - SELECT DISTINCT reaction.reaction_id - FROM ord.reaction - JOIN ord.reaction_outcome on reaction_outcome.reaction_id = reaction.id - JOIN ord.product_compound on product_compound.reaction_outcome_id = reaction_outcome.id - JOIN ord.product_measurement on product_measurement.product_compound_id = product_compound.id - JOIN ord.percentage on percentage.product_measurement_id = product_measurement.id - WHERE product_measurement.type = 'YIELD' - """ + def where_predicate(self) -> tuple[str, list]: + """Returns the query predicate and any query parameters.""" + conditions = "product_measurement.type = 'YIELD'" params = [] if self._min_yield is not None: - query += "AND percentage.value >= %s\n" + conditions += " AND percentage.value >= %s" params.append(self._min_yield) if self._max_yield is not None: - query += "AND percentage.value <= %s\n" + conditions += " AND percentage.value <= %s" params.append(self._max_yield) - return query, params + predicate = f""" + EXISTS ( + SELECT 1 FROM ord.reaction_outcome + JOIN ord.product_compound ON product_compound.reaction_outcome_id = reaction_outcome.id + JOIN ord.product_measurement ON product_measurement.product_compound_id = product_compound.id + JOIN ord.percentage ON percentage.product_measurement_id = product_measurement.id + WHERE reaction_outcome.reaction_id = reaction.id AND {conditions} + ) + """ + return predicate, params class DoiQuery(ReactionQuery): @@ -259,15 +270,16 @@ def __init__(self, dois: list[str]) -> None: self._dois = parsed_dois @property - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" - query = """ - SELECT DISTINCT reaction.reaction_id - FROM ord.reaction - JOIN ord.reaction_provenance ON reaction_provenance.reaction_id = reaction.id - WHERE reaction_provenance.doi = ANY (%s) + def where_predicate(self) -> tuple[str, list]: + """Returns the query predicate and any query parameters.""" + predicate = """ + EXISTS ( + SELECT 1 FROM ord.reaction_provenance + WHERE reaction_provenance.reaction_id = reaction.id + AND reaction_provenance.doi = ANY (%s) + ) """ - return query, [self._dois] + return predicate, [self._dois] class ReactionComponentQuery(ReactionQuery): @@ -331,15 +343,35 @@ def _mols_join(self) -> LiteralString: JOIN rdkit.mols ON rdkit.mols.id = product_compound.rdkit_mol_id """ + @property + def _mols_source(self) -> LiteralString: + """Returns the component tables and the correlation to the outer ``reaction`` row. + + The FROM list starts at the component table (not ``ord.reaction``) so it can be + used inside a correlated ``EXISTS`` against the outer ``reaction`` alias. + """ + if self._target == ReactionComponentQuery.Target.INPUT: + return """ + ord.reaction_input + JOIN ord.compound ON compound.reaction_input_id = reaction_input.id + JOIN rdkit.mols ON rdkit.mols.id = compound.rdkit_mol_id + WHERE reaction_input.reaction_id = reaction.id + """ + return """ + ord.reaction_outcome + JOIN ord.product_compound ON product_compound.reaction_outcome_id = reaction_outcome.id + JOIN rdkit.mols ON rdkit.mols.id = product_compound.rdkit_mol_id + WHERE reaction_outcome.reaction_id = reaction.id + """ + @property def is_similarity(self) -> bool: """Whether this is a similarity query, whose matches can be ranked by score.""" return self._match_mode == ReactionComponentQuery.MatchMode.SIMILAR @property - def query_and_parameters(self) -> tuple[str, list]: - """Returns the query and any query parameters.""" - mols_sql = self._mols_join + def where_predicate(self) -> tuple[str, list]: + """Returns the query predicate and any query parameters.""" if self._match_mode == ReactionComponentQuery.MatchMode.EXACT: predicate_sql = "rdkit.mols.smiles = %s" params = [Chem.CanonSmiles(self._pattern)] @@ -358,13 +390,12 @@ def query_and_parameters(self) -> tuple[str, list]: params = [self._pattern] else: raise NotImplementedError(f"Unsupported match_mode: {self._match_mode}") - query = f""" - SELECT DISTINCT reaction.reaction_id - FROM ord.reaction - {mols_sql} - WHERE {predicate_sql} + predicate = f""" + EXISTS ( + SELECT 1 FROM {self._mols_source} AND {predicate_sql} + ) """ - return query, params + return predicate, params def similarity_score_query(self) -> tuple[LiteralString, list]: """Returns a query ranking reactions by similarity, plus its leading parameter. @@ -459,11 +490,11 @@ async def run_queries( # several, results are returned unordered. ranking_query = similarity_queries[0] if len(similarity_queries) == 1 else None - queries, combined_params = [], [] + predicates, combined_params = [], [] config: dict[str, str] = {} for reaction_query in queries_list: - query, params = reaction_query.query_and_parameters - queries.append(query) + predicate, params = reaction_query.where_predicate + predicates.append(predicate) combined_params.extend(params) for name, value in reaction_query.session_config.items(): existing = config.setdefault(name, value) @@ -471,14 +502,18 @@ async def run_queries( raise ValueError( f"Conflicting values for {name}: {existing} != {value}" ) - combined_query = "\nINTERSECT\n".join(queries) + if not predicates: + raise ValueError("No query parameters were specified.") + # Each predicate is an EXISTS (or direct) condition over a single ord.reaction row, + # ANDed together. The planner can lead with the most selective one as a semi-join + # and -- crucially -- stop once LIMIT rows are found, instead of materializing one + # SELECT DISTINCT per predicate and intersecting them. + where = " AND ".join(f"({predicate})" for predicate in predicates) + combined_query = f"SELECT reaction.reaction_id FROM ord.reaction WHERE {where}" if limit and ranking_query is None: - # LIMIT sits on top of the whole INTERSECT. Each branch is materialized - # (sort + unique for DISTINCT) before the set operation, so this truncates - # the final result rather than bounding the per-branch index scans. When - # ranking, the LIMIT is deferred to the scoring step so it selects the most - # similar matches rather than an arbitrary subset. - combined_query += "LIMIT %s" + # When ranking, the LIMIT is deferred to the scoring step so it selects the + # most similar matches rather than an arbitrary subset. + combined_query += "\nLIMIT %s" combined_params.append(limit) # Apply RDKit GUCs so the substructure/similarity operators can use their GiST # indexes. These are set transaction-locally (set_config local=true), so they @@ -487,7 +522,9 @@ async def run_queries( for name, value in config.items(): await cursor.execute("SELECT set_config(%s, %s, true)", (name, value)) logger.debug((combined_query, combined_params)) - await cursor.execute(combined_query, combined_params) + # The SQL text is assembled only from internal predicate fragments (all user values + # are bound via combined_params), so treating it as a trusted query string is safe. + await cursor.execute(cast(LiteralString, combined_query), combined_params) reaction_ids = await fetch_results(cursor) if ranking_query is not None: reaction_ids = await _rank_by_similarity( diff --git a/ord_interface/api/search.py b/ord_interface/api/search.py index 1042ac6..7751510 100644 --- a/ord_interface/api/search.py +++ b/ord_interface/api/search.py @@ -68,6 +68,16 @@ BOND_LENGTH = 20 MAX_RESULTS = 1000 +# Postgres statement_timeout (ms) applied to every search connection so a pathological +# query (e.g. a common-scaffold substructure match) is cancelled and reported rather +# than running for tens of seconds. Override with ORD_INTERFACE_STATEMENT_TIMEOUT_MS. +STATEMENT_TIMEOUT_MS = int(os.getenv("ORD_INTERFACE_STATEMENT_TIMEOUT_MS", "20000")) + +QUERY_TOO_BROAD_DETAIL = ( + "The search was too broad and timed out. Add more constraints -- for example a " + "yield or conversion filter, a dataset, or a more specific structure." +) + @asynccontextmanager async def get_cursor() -> AsyncIterator[AsyncCursor[dict[str, Any]]]: @@ -85,7 +95,9 @@ async def get_cursor() -> AsyncIterator[AsyncCursor[dict[str, Any]]]: ), ) async with await psycopg.AsyncConnection[dict[str, Any]].connect( - dsn, row_factory=dict_row, options="-c search_path=public,ord" + dsn, + row_factory=dict_row, + options=f"-c search_path=public,ord -c statement_timeout={STATEMENT_TIMEOUT_MS}", ) as connection: await connection.set_read_only(True) async with connection.cursor() as cursor: @@ -168,11 +180,15 @@ async def run_query( limit = MAX_RESULTS if params.limit: limit = min(params.limit, MAX_RESULTS) - async with get_cursor() as cursor: - reaction_ids = await run_queries(cursor, queries, limit=limit) - if return_ids: - return reaction_ids - return await fetch_reactions(cursor, reaction_ids) + try: + async with get_cursor() as cursor: + reaction_ids = await run_queries(cursor, queries, limit=limit) + if return_ids: + return reaction_ids + return await fetch_reactions(cursor, reaction_ids) + except psycopg.errors.QueryCanceled as error: + # statement_timeout fired: the query was too expensive to run to completion. + raise HTTPException(status_code=400, detail=QUERY_TOO_BROAD_DETAIL) from error @router.get("/query") @@ -262,9 +278,23 @@ async def get_search_results(inputs: ReactionIdList): async def run_task(task_id: str, params: QueryParams) -> bool: - """Wraps run_query() as a background task.""" - # NOTE(skearnes): Use reaction IDs to avoid stuffing full protos into the result database. - result = await run_query(params, return_ids=True) + """Wraps run_query() as a background task. + + A failed query (e.g. a statement_timeout mapped to an HTTPException) is recorded + under ``error:{task_id}`` so ``fetch_query_result`` can report it instead of + leaving the task pending forever. + """ + try: + # NOTE(skearnes): Use reaction IDs to avoid stuffing full protos into the result database. + result = await run_query(params, return_ids=True) + except HTTPException as error: + logger.info(f"Task {task_id} failed: {error.detail}") + async with get_redis() as client: + return await client.set( + f"error:{task_id}", + json.dumps({"status_code": error.status_code, "detail": error.detail}), + ex=60 * 60, + ) logger.debug(f"Finished task {task_id}") async with get_redis() as client: return await client.set(f"result:{task_id}", json.dumps(result), ex=60 * 60) @@ -291,7 +321,11 @@ async def fetch_query_result(task_id: str): return Response( f"Task {task_id} does not exist", status_code=status.HTTP_404_NOT_FOUND ) + error = await client.get(f"error:{task_id}") result = await client.get(f"result:{task_id}") + if error is not None: + failure = json.loads(error) + return Response(failure["detail"], status_code=failure["status_code"]) if result is None: return Response( f"Task {task_id} is pending", status_code=status.HTTP_202_ACCEPTED diff --git a/ord_interface/api/search_test.py b/ord_interface/api/search_test.py index b0f7796..041d815 100644 --- a/ord_interface/api/search_test.py +++ b/ord_interface/api/search_test.py @@ -15,13 +15,18 @@ """Tests for ord_interface.api.search.""" import gzip +from contextlib import asynccontextmanager +import psycopg import pytest +from fastapi import HTTPException from ord_schema.proto import dataset_pb2 from rdkit import Chem from tenacity import retry, stop_after_attempt, wait_fixed +from ord_interface.api import search from ord_interface.api.queries import QueryResult +from ord_interface.api.search import QueryParams, run_query, run_task QUERY_PARAMS = [ # Single factor queries. @@ -177,3 +182,47 @@ def test_get_product_stats(test_client): response.raise_for_status() product_stats = response.json() assert len(product_stats) == 10 + + +@asynccontextmanager +async def _dummy_cursor(): + yield None + + +@pytest.mark.asyncio +async def test_run_query_timeout_maps_to_400(monkeypatch): + # A statement_timeout surfaces as psycopg QueryCanceled; the endpoint should + # translate it into a graceful 400 rather than letting it become a 500. + monkeypatch.setattr(search, "get_cursor", _dummy_cursor) + + async def cancel(*args, **kwargs): + raise psycopg.errors.QueryCanceled("canceling statement due to statement timeout") + + monkeypatch.setattr(search, "run_queries", cancel) + with pytest.raises(HTTPException) as excinfo: + await run_query(QueryParams(component=["c1ccccc1;input;substructure"]), return_ids=True) + assert excinfo.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_run_task_records_error(monkeypatch): + stored = {} + + class FakeRedis: + async def set(self, key, value, ex=None): + stored[key] = value + return True + + @asynccontextmanager + async def fake_redis(): + yield FakeRedis() + + async def fail(params, return_ids): + raise HTTPException(status_code=400, detail="too broad") + + monkeypatch.setattr(search, "get_redis", fake_redis) + monkeypatch.setattr(search, "run_query", fail) + await run_task("task-1", QueryParams(component=["c1ccccc1;input;substructure"])) + assert "error:task-1" in stored + assert "result:task-1" not in stored + assert '"status_code": 400' in stored["error:task-1"] From e863d251a6fee93be7f5e867f09abd3b7c1b7b21 Mon Sep 17 00:00:00 2001 From: Steven Kearnes Date: Sat, 27 Jun 2026 20:47:45 -0400 Subject: [PATCH 2/2] Apply ruff format to search_test.py CI runs `ruff format --check`; reformat the new timeout/error-path test params. Co-Authored-By: Claude Opus 4.8 (1M context) --- ord_interface/api/search_test.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ord_interface/api/search_test.py b/ord_interface/api/search_test.py index 041d815..ae0ec2a 100644 --- a/ord_interface/api/search_test.py +++ b/ord_interface/api/search_test.py @@ -196,11 +196,15 @@ async def test_run_query_timeout_maps_to_400(monkeypatch): monkeypatch.setattr(search, "get_cursor", _dummy_cursor) async def cancel(*args, **kwargs): - raise psycopg.errors.QueryCanceled("canceling statement due to statement timeout") + raise psycopg.errors.QueryCanceled( + "canceling statement due to statement timeout" + ) monkeypatch.setattr(search, "run_queries", cancel) with pytest.raises(HTTPException) as excinfo: - await run_query(QueryParams(component=["c1ccccc1;input;substructure"]), return_ids=True) + await run_query( + QueryParams(component=["c1ccccc1;input;substructure"]), return_ids=True + ) assert excinfo.value.status_code == 400