From 3bfcf79bf86c9483f203524950aab3483d5ec15c Mon Sep 17 00:00:00 2001 From: matthewberry Date: Thu, 22 Oct 2026 09:58:06 -0500 Subject: [PATCH 1/5] feat: stream full CSV search results up to CSV_MAX_ROWS Without an explicit limit, a CSV search was cut at AUTO_PAGINATION_THRESHOLD (5000) with no sign it was incomplete. CSV now streams every match through a server-side cursor, up to CSV_MAX_ROWS (default 100,000). A query that matches more, or an explicit limit above the cap, is rejected with a 400 instead of returning a partial file. CSV no longer runs the COUNT query, whose result it never used. Closes #18 --- app/core/config.py | 2 + app/db/database.py | 11 +++++- app/db/queries.py | 41 ++++++++++++++----- app/main.py | 5 +++ app/routers/search.py | 85 +++++++++++++++++++++++++--------------- tests/test_search_csv.py | 79 ++++++++++++++++++++++++++++++++----- 6 files changed, 172 insertions(+), 51 deletions(-) diff --git a/app/core/config.py b/app/core/config.py index 8c0415c..0f90187 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -12,6 +12,8 @@ class Settings(BaseSettings): # API behavior configuration AUTO_PAGINATION_THRESHOLD: int = 5000 + # Most records a single CSV search response may contain + CSV_MAX_ROWS: int = 100_000 # CORS configuration CORS_ORIGINS: List[str] = ["*"] diff --git a/app/db/database.py b/app/db/database.py index dc53593..5150ebd 100644 --- a/app/db/database.py +++ b/app/db/database.py @@ -1,4 +1,4 @@ -from typing import Any, AsyncGenerator, Dict, List, Optional +from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional import asyncpg from loguru import logger @@ -71,6 +71,15 @@ async def fetch(self, query: str, *args, **kwargs) -> List[Dict[str, Any]]: records = await self.pool.fetch(query, *args, **kwargs) return [dict(record) for record in records] + async def iterate(self, query: str, *args) -> AsyncIterator[Dict[str, Any]]: + """Stream rows from a query through a server-side cursor.""" + if not self.pool: + await self.connect() + # asyncpg cursors only work inside a transaction + async with self.pool.acquire() as conn, conn.transaction(): + async for record in conn.cursor(query, *args, prefetch=1000): + yield dict(record) + async def fetchval(self, query: str, *args, **kwargs) -> Any: """Fetch a single value from a query.""" if not self.pool: diff --git a/app/db/queries.py b/app/db/queries.py index 5e78398..fbf6c86 100644 --- a/app/db/queries.py +++ b/app/db/queries.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Tuple +from typing import Any, AsyncIterator, Dict, List, Tuple import rich import re @@ -138,10 +138,8 @@ def get_query(columns_to_select: str, where_clause: str, include_order_by: bool ORDER BY {parse_ordering(ordering)}""" return query -async def get_filtered_data( - db: Database, params: CLEANSearchQueryParams -) -> List[Dict[str, Any]]: - """Get filtered data from the database.""" +async def build_data_query(params: CLEANSearchQueryParams) -> Tuple[str, List[Any]]: + """Build the filtered data query and its arguments.""" where_clause, query_params = await build_conditions(params) columns_to_select = """ @@ -173,11 +171,36 @@ async def get_filtered_data( rich.print(f"{query=}") # Extract query parameters from the dictionary - query_args = list(query_params.values()) + return query, list(query_params.values()) - # Execute the query - records = await db.fetch(query, *query_args) - return records + +async def get_filtered_data( + db: Database, params: CLEANSearchQueryParams +) -> List[Dict[str, Any]]: + """Get filtered data from the database.""" + query, query_args = await build_data_query(params) + return await db.fetch(query, *query_args) + + +async def iter_filtered_data( + db: Database, params: CLEANSearchQueryParams +) -> AsyncIterator[Dict[str, Any]]: + """Stream filtered data from the database.""" + query, query_args = await build_data_query(params) + async for record in db.iterate(query, *query_args): + yield record + + +async def has_more_rows_than(db: Database, params: CLEANSearchQueryParams, n: int) -> bool: + """Check whether more than n records match the filters, starting at params.offset.""" + where_clause, query_params = await build_conditions(params) + + # Stops at the first row past n instead of counting every match + query = f"""SELECT EXISTS ( + {get_query("1", where_clause, include_order_by=False)} + OFFSET {(params.offset or 0) + n})""" + + return await db.fetchval(query, *query_params.values()) async def get_total_count(db: Database, params: CLEANSearchQueryParams) -> int: diff --git a/app/main.py b/app/main.py index cf72447..8b55fd9 100644 --- a/app/main.py +++ b/app/main.py @@ -57,6 +57,11 @@ async def lifespan(app: FastAPI): This threshold can be configured using the `AUTO_PAGINATION_THRESHOLD` environment variable. +CSV responses are not paginated. Without an explicit limit, a CSV search returns every matching +record, up to {settings.CSV_MAX_ROWS}; a query that matches more is rejected with a 400 asking +you to narrow the filters or page with `limit` and `offset`. This maximum can be configured using +the `CSV_MAX_ROWS` environment variable. + ## Data License The CLEAN Database dataset is licensed under the Creative Commons Attribution 4.0 International diff --git a/app/routers/search.py b/app/routers/search.py index d0b0ff9..c215680 100644 --- a/app/routers/search.py +++ b/app/routers/search.py @@ -1,6 +1,6 @@ import csv from io import StringIO -from typing import Any, List, Literal, Optional +from typing import Any, AsyncIterator, Dict, List, Literal, Optional from urllib.parse import urlencode from fastapi import APIRouter, Depends, HTTPException, Query, Request @@ -9,13 +9,32 @@ from app.core.config import settings from app.db.database import Database, get_db -from app.db.queries import get_ec_suggestions, get_filtered_data, get_total_count, get_typeahead_suggestions +from app.db.queries import get_ec_suggestions, get_filtered_data, get_total_count, get_typeahead_suggestions, has_more_rows_than, iter_filtered_data from app.models.query_params import CLEANECLookupQueryParams, CLEANSearchQueryParams, CLEANTypeaheadQueryParams, ResponseFormat from app.models.clean_data import CLEANColumn, CLEANDataBase, CLEANECLookupResponse, CLEANECLookupMatch, CLEANSearchResponse, CLEANTypeaheadResponse, CurationStatusOption, CLEANCurationStatusResponse router = APIRouter(tags=["Search"]) +async def stream_csv(rows: AsyncIterator[Dict[str, Any]]) -> AsyncIterator[str]: + """Write rows as CSV, one chunk per row.""" + output = StringIO() + writer = csv.DictWriter( + output, fieldnames=[c.value for c in CLEANColumn], extrasaction="ignore" + ) + writer.writeheader() + async for row in rows: + # Join array columns into one ';'-separated cell + writer.writerow({ + k: ";".join(str(v) for v in value) if isinstance(value, list) else value + for k, value in row.items() + }) + yield output.getvalue() + output.seek(0) + output.truncate() + yield output.getvalue() + + def parse_query_params( # String filters accession: Optional[List[str]] = Query( @@ -114,8 +133,10 @@ async def get_data( Filters that accept multiple values on the same parameter (e.g. `organism`) are combined with OR logic, while filters on different parameters are combined with AND logic. -The response format can be either JSON (default) or CSV. Results are automatically -paginated when no explicit `limit` is provided. +The response format can be either JSON (default) or CSV. JSON results are automatically +paginated when no explicit `limit` is provided. CSV results are not paginated: without a +`limit`, every matching record is returned, up to `CSV_MAX_ROWS` (100,000 by default); a +query that matches more is rejected with a 400. ### URL examples @@ -204,41 +225,39 @@ async def get_data( """ try: - # Apply default page size if no explicit limit provided - if params.limit is None: - params.limit = settings.AUTO_PAGINATION_THRESHOLD - - # Get total count for the query (without pagination) - total_count = await get_total_count(db, params) - - # Get data from database - data = await get_filtered_data(db, params) - - # Handle response format + # CSV is not paginated: without an explicit limit, return every match up to the cap if params.format == ResponseFormat.CSV: - # Create CSV response - output = StringIO() + max_rows = settings.CSV_MAX_ROWS + if params.limit is None: + params.limit = max_rows + if await has_more_rows_than(db, params, max_rows): + raise HTTPException( + status_code=400, + detail=f"Query matches more than {max_rows} records, the most a CSV " + "response can contain. Narrow the filters, or page with limit and offset.", + ) + elif params.limit > max_rows: + raise HTTPException( + status_code=400, + detail=f"limit cannot exceed {max_rows} for CSV responses", + ) - # Write CSV - writer = csv.DictWriter( - output, fieldnames=[c.value for c in CLEANColumn], extrasaction="ignore" - ) - writer.writeheader() - for row in data: - # Join array columns into one ';'-separated cell - writer.writerow({ - k: ";".join(str(v) for v in value) if isinstance(value, list) else value - for k, value in row.items() - }) - - # Return streaming response - output.seek(0) return StreamingResponse( - iter([output.getvalue()]), + stream_csv(iter_filtered_data(db, params)), media_type="text/csv", headers={"Content-Disposition": "attachment; filename=CLEAN_data.csv"}, ) else: + # Apply default page size if no explicit limit provided + if params.limit is None: + params.limit = settings.AUTO_PAGINATION_THRESHOLD + + # Get total count for the query (without pagination) + total_count = await get_total_count(db, params) + + # Get data from database + data = await get_filtered_data(db, params) + response = CLEANSearchResponse( total=total_count, offset=params.offset, @@ -311,6 +330,8 @@ async def get_data( return response + except HTTPException: + raise except Exception as e: logger.error(f"Error getting data: {e}") raise HTTPException(status_code=500, detail=f"Error retrieving data: {str(e)}") diff --git a/tests/test_search_csv.py b/tests/test_search_csv.py index 8fd75ae..01e0506 100644 --- a/tests/test_search_csv.py +++ b/tests/test_search_csv.py @@ -6,6 +6,7 @@ from fastapi.testclient import TestClient import app.routers.search as search +from app.core.config import settings from app.db.database import get_db from app.main import app from app.models.clean_data import CLEANColumn @@ -20,12 +21,33 @@ def client(): app.dependency_overrides.clear() -def get_csv(client, rows): +@pytest.fixture +def db_calls(): + """Mock the DB query functions; iter_filtered_data yields the rows in db_calls.rows.""" + class Calls: + rows = [] + streamed_params = [] + + async def iter_filtered_data(db, params): + Calls.streamed_params.append(params.model_copy()) + for row in Calls.rows: + yield row + with ( - patch.object(search, "get_total_count", AsyncMock(return_value=len(rows))), - patch.object(search, "get_filtered_data", AsyncMock(return_value=rows)), + patch.object(search, "iter_filtered_data", iter_filtered_data), + patch.object(search, "has_more_rows_than", AsyncMock(return_value=False)) as more, + patch.object(search, "get_total_count", AsyncMock()) as count, + patch.object(search, "get_filtered_data", AsyncMock()) as fetch, ): - response = client.get("/api/v1/search?accession=P51451&format=csv&limit=10") + Calls.has_more_rows_than, Calls.get_total_count, Calls.get_filtered_data = more, count, fetch + yield Calls + + +def get_csv(client, query="accession=P51451"): + return client.get(f"/api/v1/search?{query}&format=csv") + + +def parse(response): assert response.status_code == 200, response.text return list(csv.reader(StringIO(response.text))) @@ -37,12 +59,13 @@ def make_row(**overrides): return row -def test_csv_joins_array_columns(client): - rows = get_csv(client, [make_row( +def test_csv_joins_array_columns(client, db_calls): + db_calls.rows = [make_row( clean_ec_number_array=["2.7.10.2", "2.7.11.1"], clean_ec_confidence_array=[0.98, 0.12], annot_ec_number_array=None, - )]) + )] + rows = parse(get_csv(client, "accession=P51451&limit=10")) assert rows[0] == HEADER record = dict(zip(HEADER, rows[1])) @@ -52,5 +75,43 @@ def test_csv_joins_array_columns(client): assert record["annot_ec_number_array"] == "" -def test_csv_has_header_when_no_rows_match(client): - assert get_csv(client, []) == [HEADER] +def test_csv_has_header_when_no_rows_match(client, db_calls): + assert parse(get_csv(client)) == [HEADER] + + +def test_csv_without_limit_returns_all_rows_up_to_cap(client, db_calls): + db_calls.rows = [make_row(accession=f"P{i:05d}") for i in range(settings.AUTO_PAGINATION_THRESHOLD + 1)] + rows = parse(get_csv(client, "offset=7")) + + assert len(rows) == 1 + settings.AUTO_PAGINATION_THRESHOLD + 1 + assert db_calls.streamed_params[0].limit == settings.CSV_MAX_ROWS + assert db_calls.streamed_params[0].offset == 7 + db_calls.has_more_rows_than.assert_awaited_once() + assert db_calls.has_more_rows_than.await_args.args[2] == settings.CSV_MAX_ROWS + db_calls.get_total_count.assert_not_awaited() + db_calls.get_filtered_data.assert_not_awaited() + + +def test_csv_without_limit_rejects_more_rows_than_cap(client, db_calls): + db_calls.has_more_rows_than.return_value = True + response = get_csv(client) + + assert response.status_code == 400 + assert f"more than {settings.CSV_MAX_ROWS} records" in response.json()["detail"] + assert db_calls.streamed_params == [] + + +def test_csv_with_limit_skips_cap_check(client, db_calls): + parse(get_csv(client, "limit=50&offset=100")) + + assert db_calls.streamed_params[0].limit == 50 + assert db_calls.streamed_params[0].offset == 100 + db_calls.has_more_rows_than.assert_not_awaited() + + +def test_csv_rejects_limit_above_cap(client, db_calls): + response = get_csv(client, f"limit={settings.CSV_MAX_ROWS + 1}") + + assert response.status_code == 400 + assert response.json()["detail"] == f"limit cannot exceed {settings.CSV_MAX_ROWS} for CSV responses" + assert db_calls.streamed_params == [] From c990f0232f33b30df5e79aa107e2c68d62cd6239 Mon Sep 17 00:00:00 2001 From: matthewberry Date: Thu, 22 Oct 2026 10:08:17 -0500 Subject: [PATCH 2/5] fix: spool CSV to a temp file before sending Streaming straight from the cursor held a pooled DB connection, inside a transaction, for as long as the client took to read the response. Ten slow clients exhausted the pool and stalled every endpoint. Write the CSV to a SpooledTemporaryFile (in memory up to 8 MB, then on disk) and release the connection before sending. A DB error now returns a 500 instead of a truncated file. The file is closed in a BackgroundTask, which also runs when the client disconnects. --- app/routers/search.py | 36 ++++++++++++++++++++++++++---------- tests/test_search_csv.py | 23 +++++++++++++++++++++++ 2 files changed, 49 insertions(+), 10 deletions(-) diff --git a/app/routers/search.py b/app/routers/search.py index c215680..8eed3bf 100644 --- a/app/routers/search.py +++ b/app/routers/search.py @@ -1,10 +1,11 @@ import csv -from io import StringIO -from typing import Any, AsyncIterator, Dict, List, Literal, Optional +from tempfile import SpooledTemporaryFile +from typing import Any, AsyncIterator, Dict, Iterator, List, Literal, Optional from urllib.parse import urlencode from fastapi import APIRouter, Depends, HTTPException, Query, Request from fastapi.responses import StreamingResponse +from starlette.background import BackgroundTask from loguru import logger from app.core.config import settings @@ -16,9 +17,17 @@ router = APIRouter(tags=["Search"]) -async def stream_csv(rows: AsyncIterator[Dict[str, Any]]) -> AsyncIterator[str]: - """Write rows as CSV, one chunk per row.""" - output = StringIO() +# CSV output stays in memory up to this size, then spills to a temp file on disk +CSV_SPOOL_MAX_MEMORY = 8 * 1024 * 1024 + + +async def write_csv(rows: AsyncIterator[Dict[str, Any]]) -> SpooledTemporaryFile: + """Write rows as CSV to a spooled temp file, rewound for reading. + + The whole result is written before the response starts, so the DB connection is + released at the speed of the DB, not the speed of the client reading the response. + """ + output = SpooledTemporaryFile(max_size=CSV_SPOOL_MAX_MEMORY, mode="w+", newline="") writer = csv.DictWriter( output, fieldnames=[c.value for c in CLEANColumn], extrasaction="ignore" ) @@ -29,10 +38,14 @@ async def stream_csv(rows: AsyncIterator[Dict[str, Any]]) -> AsyncIterator[str]: k: ";".join(str(v) for v in value) if isinstance(value, list) else value for k, value in row.items() }) - yield output.getvalue() - output.seek(0) - output.truncate() - yield output.getvalue() + output.seek(0) + return output + + +def read_chunks(file: SpooledTemporaryFile, size: int = 64 * 1024) -> Iterator[str]: + """Read a file in chunks.""" + while chunk := file.read(size): + yield chunk def parse_query_params( @@ -242,10 +255,13 @@ async def get_data( detail=f"limit cannot exceed {max_rows} for CSV responses", ) + csv_file = await write_csv(iter_filtered_data(db, params)) return StreamingResponse( - stream_csv(iter_filtered_data(db, params)), + read_chunks(csv_file), media_type="text/csv", headers={"Content-Disposition": "attachment; filename=CLEAN_data.csv"}, + # Runs after the response ends, including when the client disconnects + background=BackgroundTask(csv_file.close), ) else: # Apply default page size if no explicit limit provided diff --git a/tests/test_search_csv.py b/tests/test_search_csv.py index 01e0506..abe713d 100644 --- a/tests/test_search_csv.py +++ b/tests/test_search_csv.py @@ -109,6 +109,29 @@ def test_csv_with_limit_skips_cap_check(client, db_calls): db_calls.has_more_rows_than.assert_not_awaited() +def test_csv_db_error_mid_result_returns_500(client): + async def failing_rows(db, params): + yield make_row() + raise RuntimeError("connection lost") + + with ( + patch.object(search, "iter_filtered_data", failing_rows), + patch.object(search, "has_more_rows_than", AsyncMock(return_value=False)), + ): + response = get_csv(client) + + assert response.status_code == 500 + assert "connection lost" in response.json()["detail"] + + +def test_csv_spills_large_output_to_disk(client, db_calls): + db_calls.rows = [make_row(protein_sequence="M" * 100_000) for _ in range(100)] + rows = parse(get_csv(client)) + + assert len(rows) == 101 + assert all(len(dict(zip(HEADER, r))["protein_sequence"]) == 100_000 for r in rows[1:]) + + def test_csv_rejects_limit_above_cap(client, db_calls): response = get_csv(client, f"limit={settings.CSV_MAX_ROWS + 1}") From 0733f0ac1438ed98adc59ebf00f2e7417004ffec Mon Sep 17 00:00:00 2001 From: matthewberry Date: Thu, 22 Oct 2026 10:58:53 -0500 Subject: [PATCH 3/5] refactor: use one 100,000 search cap for limit and CSV Replace the configurable CSV_MAX_ROWS with MAX_SEARCH_LIMIT, the same bound #25 put on limit. An explicit limit above it already gets a 422, so drop the CSV-only 400 for that case and its test. Clamp the cap check's OFFSET to MAX_BIGINT: with offset allowed up to the bigint maximum, offset + cap overflowed and Postgres returned a 500. --- app/core/config.py | 2 -- app/db/queries.py | 9 +++++++-- app/main.py | 5 ++--- app/routers/search.py | 24 +++++++++--------------- tests/test_search_csv.py | 18 ++++++++++-------- 5 files changed, 28 insertions(+), 30 deletions(-) diff --git a/app/core/config.py b/app/core/config.py index 0f90187..8c0415c 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -12,8 +12,6 @@ class Settings(BaseSettings): # API behavior configuration AUTO_PAGINATION_THRESHOLD: int = 5000 - # Most records a single CSV search response may contain - CSV_MAX_ROWS: int = 100_000 # CORS configuration CORS_ORIGINS: List[str] = ["*"] diff --git a/app/db/queries.py b/app/db/queries.py index 926f762..34faf61 100644 --- a/app/db/queries.py +++ b/app/db/queries.py @@ -5,6 +5,9 @@ from app.models.clean_data import CLEANColumn from app.models.query_params import CLEANECLookupQueryParams, CLEANSearchQueryParams, CLEANTypeaheadQueryParams +# LIMIT and OFFSET are bigint in Postgres; larger values fail when the query runs. +MAX_BIGINT = 2**63 - 1 + async def build_conditions( params: CLEANSearchQueryParams, ) -> Tuple[str, Dict[str, Any]]: @@ -192,10 +195,12 @@ async def has_more_rows_than(db: Database, params: CLEANSearchQueryParams, n: in """Check whether more than n records match the filters, starting at params.offset.""" where_clause, query_params = await build_conditions(params) - # Stops at the first row past n instead of counting every match + # Stops at the first row past n instead of counting every match. Clamped so OFFSET + # stays a bigint; no table has MAX_BIGINT rows, so the answer is still correct. + offset = min((params.offset or 0) + n, MAX_BIGINT) query = f"""SELECT EXISTS ( {get_query("1", where_clause, include_order_by=False)} - OFFSET {(params.offset or 0) + n})""" + OFFSET {offset})""" return await db.fetchval(query, *query_params.values()) diff --git a/app/main.py b/app/main.py index 8b55fd9..18ef53c 100644 --- a/app/main.py +++ b/app/main.py @@ -58,9 +58,8 @@ async def lifespan(app: FastAPI): This threshold can be configured using the `AUTO_PAGINATION_THRESHOLD` environment variable. CSV responses are not paginated. Without an explicit limit, a CSV search returns every matching -record, up to {settings.CSV_MAX_ROWS}; a query that matches more is rejected with a 400 asking -you to narrow the filters or page with `limit` and `offset`. This maximum can be configured using -the `CSV_MAX_ROWS` environment variable. +record, up to {search.MAX_SEARCH_LIMIT}; a query that matches more is rejected with a 400 asking +you to narrow the filters or page with `limit` and `offset`. ## Data License diff --git a/app/routers/search.py b/app/routers/search.py index 8049f7a..e8b0101 100644 --- a/app/routers/search.py +++ b/app/routers/search.py @@ -10,14 +10,14 @@ from app.core.config import settings from app.db.database import Database, get_db -from app.db.queries import get_ec_suggestions, get_filtered_data, get_total_count, get_typeahead_suggestions, has_more_rows_than, iter_filtered_data +from app.db.queries import MAX_BIGINT, get_ec_suggestions, get_filtered_data, get_total_count, get_typeahead_suggestions, has_more_rows_than, iter_filtered_data from app.models.query_params import CLEANECLookupQueryParams, CLEANSearchQueryParams, CLEANTypeaheadQueryParams, ResponseFormat from app.models.clean_data import CLEANColumn, CLEANDataBase, CLEANECLookupResponse, CLEANECLookupMatch, CLEANSearchResponse, CLEANTypeaheadResponse, CurationStatusOption, CLEANCurationStatusResponse router = APIRouter(tags=["Search"]) -# LIMIT and OFFSET are bigint in Postgres; larger values fail when the query runs. -MAX_BIGINT = 2**63 - 1 +# Most records one /search response may contain, as an explicit limit or a whole CSV +MAX_SEARCH_LIMIT = 100_000 # CSV output stays in memory up to this size, then spills to a temp file on disk @@ -96,7 +96,7 @@ def parse_query_params( default=ResponseFormat.JSON, description="Response format (json or csv)" ), limit: Optional[int] = Query( - None, ge=1, le=100000, description="Maximum number of records to return" + None, ge=1, le=MAX_SEARCH_LIMIT, description="Maximum number of records to return" ), offset: Optional[int] = Query(0, ge=0, le=MAX_BIGINT, description="Number of records to skip"), ordering: Optional[str] = Query( @@ -151,8 +151,8 @@ async def get_data( The response format can be either JSON (default) or CSV. JSON results are automatically paginated when no explicit `limit` is provided. CSV results are not paginated: without a -`limit`, every matching record is returned, up to `CSV_MAX_ROWS` (100,000 by default); a -query that matches more is rejected with a 400. +`limit`, every matching record is returned, up to 100,000; a query that matches more is +rejected with a 400. ### URL examples @@ -243,20 +243,14 @@ async def get_data( try: # CSV is not paginated: without an explicit limit, return every match up to the cap if params.format == ResponseFormat.CSV: - max_rows = settings.CSV_MAX_ROWS if params.limit is None: - params.limit = max_rows - if await has_more_rows_than(db, params, max_rows): + params.limit = MAX_SEARCH_LIMIT + if await has_more_rows_than(db, params, MAX_SEARCH_LIMIT): raise HTTPException( status_code=400, - detail=f"Query matches more than {max_rows} records, the most a CSV " + detail=f"Query matches more than {MAX_SEARCH_LIMIT} records, the most a CSV " "response can contain. Narrow the filters, or page with limit and offset.", ) - elif params.limit > max_rows: - raise HTTPException( - status_code=400, - detail=f"limit cannot exceed {max_rows} for CSV responses", - ) csv_file = await write_csv(iter_filtered_data(db, params)) return StreamingResponse( diff --git a/tests/test_search_csv.py b/tests/test_search_csv.py index 3b69697..7519647 100644 --- a/tests/test_search_csv.py +++ b/tests/test_search_csv.py @@ -6,6 +6,7 @@ import app.routers.search as search from app.core.config import settings +from app.db.queries import MAX_BIGINT from app.models.clean_data import CLEANColumn HEADER = [c.value for c in CLEANColumn] @@ -74,10 +75,10 @@ def test_csv_without_limit_returns_all_rows_up_to_cap(client, db_calls): rows = parse(get_csv(client, "offset=7")) assert len(rows) == 1 + settings.AUTO_PAGINATION_THRESHOLD + 1 - assert db_calls.streamed_params[0].limit == settings.CSV_MAX_ROWS + assert db_calls.streamed_params[0].limit == search.MAX_SEARCH_LIMIT assert db_calls.streamed_params[0].offset == 7 db_calls.has_more_rows_than.assert_awaited_once() - assert db_calls.has_more_rows_than.await_args.args[2] == settings.CSV_MAX_ROWS + assert db_calls.has_more_rows_than.await_args.args[2] == search.MAX_SEARCH_LIMIT db_calls.get_total_count.assert_not_awaited() db_calls.get_filtered_data.assert_not_awaited() @@ -87,7 +88,7 @@ def test_csv_without_limit_rejects_more_rows_than_cap(client, db_calls): response = get_csv(client) assert response.status_code == 400 - assert f"more than {settings.CSV_MAX_ROWS} records" in response.json()["detail"] + assert f"more than {search.MAX_SEARCH_LIMIT} records" in response.json()["detail"] assert db_calls.streamed_params == [] @@ -122,9 +123,10 @@ def test_csv_spills_large_output_to_disk(client, db_calls): assert all(len(dict(zip(HEADER, r))["protein_sequence"]) == 100_000 for r in rows[1:]) -def test_csv_rejects_limit_above_cap(client, db_calls): - response = get_csv(client, f"limit={settings.CSV_MAX_ROWS + 1}") +async def test_cap_check_offset_stays_within_bigint(db): + from app.db.queries import has_more_rows_than + from app.models.query_params import CLEANSearchQueryParams - assert response.status_code == 400 - assert response.json()["detail"] == f"limit cannot exceed {settings.CSV_MAX_ROWS} for CSV responses" - assert db_calls.streamed_params == [] + await has_more_rows_than(db, CLEANSearchQueryParams(offset=MAX_BIGINT - 1), search.MAX_SEARCH_LIMIT) + + assert db.fetchval.call_args.args[0].rstrip().endswith(f"OFFSET {MAX_BIGINT})") From 12eb4f06d17ed493ce9adc6ad527ece5f1b5a38a Mon Sep 17 00:00:00 2001 From: matthewberry Date: Thu, 22 Oct 2026 11:57:51 -0500 Subject: [PATCH 4/5] fix: enforce the CSV cap while writing instead of with a separate query The cap check and the export were two queries, so data changing between them could let the export's LIMIT silently drop rows, and every result under the cap was scanned twice. Fetch cap + 1 rows and return the 400 from write_csv if the extra row arrives. This removes has_more_rows_than and its OFFSET clamp, and leaves the caller's params unchanged. Close the row generators explicitly with aclosing, so stopping early releases the cursor's connection immediately rather than whenever the garbage collector runs. --- app/db/queries.py | 24 +++-------- app/routers/search.py | 60 +++++++++++++++++---------- tests/test_search_csv.py | 89 ++++++++++++++++++++++++++++------------ 3 files changed, 107 insertions(+), 66 deletions(-) diff --git a/app/db/queries.py b/app/db/queries.py index 34faf61..8c2ff25 100644 --- a/app/db/queries.py +++ b/app/db/queries.py @@ -1,3 +1,4 @@ +from contextlib import aclosing from typing import Any, AsyncIterator, Dict, List, Tuple import re @@ -5,9 +6,6 @@ from app.models.clean_data import CLEANColumn from app.models.query_params import CLEANECLookupQueryParams, CLEANSearchQueryParams, CLEANTypeaheadQueryParams -# LIMIT and OFFSET are bigint in Postgres; larger values fail when the query runs. -MAX_BIGINT = 2**63 - 1 - async def build_conditions( params: CLEANSearchQueryParams, ) -> Tuple[str, Dict[str, Any]]: @@ -187,22 +185,10 @@ async def iter_filtered_data( ) -> AsyncIterator[Dict[str, Any]]: """Stream filtered data from the database.""" query, query_args = await build_data_query(params) - async for record in db.iterate(query, *query_args): - yield record - - -async def has_more_rows_than(db: Database, params: CLEANSearchQueryParams, n: int) -> bool: - """Check whether more than n records match the filters, starting at params.offset.""" - where_clause, query_params = await build_conditions(params) - - # Stops at the first row past n instead of counting every match. Clamped so OFFSET - # stays a bigint; no table has MAX_BIGINT rows, so the answer is still correct. - offset = min((params.offset or 0) + n, MAX_BIGINT) - query = f"""SELECT EXISTS ( - {get_query("1", where_clause, include_order_by=False)} - OFFSET {offset})""" - - return await db.fetchval(query, *query_params.values()) + # Closing this generator early closes the cursor too, releasing its connection + async with aclosing(db.iterate(query, *query_args)) as records: + async for record in records: + yield record async def get_total_count(db: Database, params: CLEANSearchQueryParams) -> int: diff --git a/app/routers/search.py b/app/routers/search.py index e8b0101..c50d78d 100644 --- a/app/routers/search.py +++ b/app/routers/search.py @@ -1,4 +1,5 @@ import csv +from contextlib import aclosing from tempfile import SpooledTemporaryFile from typing import Any, AsyncIterator, Dict, Iterator, List, Literal, Optional from urllib.parse import urlencode @@ -10,12 +11,15 @@ from app.core.config import settings from app.db.database import Database, get_db -from app.db.queries import MAX_BIGINT, get_ec_suggestions, get_filtered_data, get_total_count, get_typeahead_suggestions, has_more_rows_than, iter_filtered_data +from app.db.queries import get_ec_suggestions, get_filtered_data, get_total_count, get_typeahead_suggestions, iter_filtered_data from app.models.query_params import CLEANECLookupQueryParams, CLEANSearchQueryParams, CLEANTypeaheadQueryParams, ResponseFormat from app.models.clean_data import CLEANColumn, CLEANDataBase, CLEANECLookupResponse, CLEANECLookupMatch, CLEANSearchResponse, CLEANTypeaheadResponse, CurationStatusOption, CLEANCurationStatusResponse router = APIRouter(tags=["Search"]) +# LIMIT and OFFSET are bigint in Postgres; larger values fail when the query runs. +MAX_BIGINT = 2**63 - 1 + # Most records one /search response may contain, as an explicit limit or a whole CSV MAX_SEARCH_LIMIT = 100_000 @@ -24,23 +28,40 @@ CSV_SPOOL_MAX_MEMORY = 8 * 1024 * 1024 -async def write_csv(rows: AsyncIterator[Dict[str, Any]]) -> SpooledTemporaryFile: +async def write_csv( + rows: AsyncIterator[Dict[str, Any]], max_rows: Optional[int] = None +) -> SpooledTemporaryFile: """Write rows as CSV to a spooled temp file, rewound for reading. The whole result is written before the response starts, so the DB connection is released at the speed of the DB, not the speed of the client reading the response. + Raises a 400 if rows has more than max_rows rows. """ output = SpooledTemporaryFile(max_size=CSV_SPOOL_MAX_MEMORY, mode="w+", newline="") - writer = csv.DictWriter( - output, fieldnames=[c.value for c in CLEANColumn], extrasaction="ignore" - ) - writer.writeheader() - async for row in rows: - # Join array columns into one ';'-separated cell - writer.writerow({ - k: ";".join(str(v) for v in value) if isinstance(value, list) else value - for k, value in row.items() - }) + try: + writer = csv.DictWriter( + output, fieldnames=[c.value for c in CLEANColumn], extrasaction="ignore" + ) + writer.writeheader() + # aclosing releases the DB connection now if we stop early, not at garbage collection + async with aclosing(rows): + count = 0 + async for row in rows: + count += 1 + if max_rows is not None and count > max_rows: + raise HTTPException( + status_code=400, + detail=f"Query matches more than {max_rows} records, the most a CSV " + "response can contain. Narrow the filters, or page with limit and offset.", + ) + # Join array columns into one ';'-separated cell + writer.writerow({ + k: ";".join(str(v) for v in value) if isinstance(value, list) else value + for k, value in row.items() + }) + except BaseException: + output.close() + raise output.seek(0) return output @@ -244,15 +265,12 @@ async def get_data( # CSV is not paginated: without an explicit limit, return every match up to the cap if params.format == ResponseFormat.CSV: if params.limit is None: - params.limit = MAX_SEARCH_LIMIT - if await has_more_rows_than(db, params, MAX_SEARCH_LIMIT): - raise HTTPException( - status_code=400, - detail=f"Query matches more than {MAX_SEARCH_LIMIT} records, the most a CSV " - "response can contain. Narrow the filters, or page with limit and offset.", - ) - - csv_file = await write_csv(iter_filtered_data(db, params)) + # Fetch one row past the cap, so a result that is too large is detected + # while writing instead of silently cut off + rows = iter_filtered_data(db, params.model_copy(update={"limit": MAX_SEARCH_LIMIT + 1})) + csv_file = await write_csv(rows, max_rows=MAX_SEARCH_LIMIT) + else: + csv_file = await write_csv(iter_filtered_data(db, params)) return StreamingResponse( read_chunks(csv_file), media_type="text/csv", diff --git a/tests/test_search_csv.py b/tests/test_search_csv.py index 7519647..a86f4f4 100644 --- a/tests/test_search_csv.py +++ b/tests/test_search_csv.py @@ -6,7 +6,6 @@ import app.routers.search as search from app.core.config import settings -from app.db.queries import MAX_BIGINT from app.models.clean_data import CLEANColumn HEADER = [c.value for c in CLEANColumn] @@ -18,19 +17,22 @@ def db_calls(): class Calls: rows = [] streamed_params = [] + rows_closed = False async def iter_filtered_data(db, params): Calls.streamed_params.append(params.model_copy()) - for row in Calls.rows: - yield row + try: + for row in Calls.rows: + yield row + finally: + Calls.rows_closed = True with ( patch.object(search, "iter_filtered_data", iter_filtered_data), - patch.object(search, "has_more_rows_than", AsyncMock(return_value=False)) as more, patch.object(search, "get_total_count", AsyncMock()) as count, patch.object(search, "get_filtered_data", AsyncMock()) as fetch, ): - Calls.has_more_rows_than, Calls.get_total_count, Calls.get_filtered_data = more, count, fetch + Calls.get_total_count, Calls.get_filtered_data = count, fetch yield Calls @@ -70,34 +72,38 @@ def test_csv_has_header_when_no_rows_match(client, db_calls): assert parse(get_csv(client)) == [HEADER] -def test_csv_without_limit_returns_all_rows_up_to_cap(client, db_calls): - db_calls.rows = [make_row(accession=f"P{i:05d}") for i in range(settings.AUTO_PAGINATION_THRESHOLD + 1)] +def test_csv_without_limit_returns_all_rows_up_to_cap(client, db_calls, monkeypatch): + monkeypatch.setattr(search, "MAX_SEARCH_LIMIT", 3) + db_calls.rows = [make_row(accession=f"P{i:05d}") for i in range(3)] rows = parse(get_csv(client, "offset=7")) - assert len(rows) == 1 + settings.AUTO_PAGINATION_THRESHOLD + 1 - assert db_calls.streamed_params[0].limit == search.MAX_SEARCH_LIMIT + assert [r[3] for r in rows[1:]] == ["P00000", "P00001", "P00002"] + # One row past the cap is requested, so an oversized result can be detected + assert db_calls.streamed_params[0].limit == 4 assert db_calls.streamed_params[0].offset == 7 - db_calls.has_more_rows_than.assert_awaited_once() - assert db_calls.has_more_rows_than.await_args.args[2] == search.MAX_SEARCH_LIMIT db_calls.get_total_count.assert_not_awaited() db_calls.get_filtered_data.assert_not_awaited() -def test_csv_without_limit_rejects_more_rows_than_cap(client, db_calls): - db_calls.has_more_rows_than.return_value = True +def test_csv_without_limit_rejects_more_rows_than_cap(client, db_calls, monkeypatch): + monkeypatch.setattr(search, "MAX_SEARCH_LIMIT", 3) + db_calls.rows = [make_row() for _ in range(4)] response = get_csv(client) assert response.status_code == 400 - assert f"more than {search.MAX_SEARCH_LIMIT} records" in response.json()["detail"] - assert db_calls.streamed_params == [] + assert response.json()["detail"].startswith("Query matches more than 3 records") + # The row iterator, and so the DB cursor, is closed as soon as the cap is passed + assert db_calls.rows_closed -def test_csv_with_limit_skips_cap_check(client, db_calls): - parse(get_csv(client, "limit=50&offset=100")) +def test_csv_with_limit_is_not_capped(client, db_calls, monkeypatch): + monkeypatch.setattr(search, "MAX_SEARCH_LIMIT", 3) + db_calls.rows = [make_row() for _ in range(4)] + rows = parse(get_csv(client, "limit=50&offset=100")) + assert len(rows) == 5 assert db_calls.streamed_params[0].limit == 50 assert db_calls.streamed_params[0].offset == 100 - db_calls.has_more_rows_than.assert_not_awaited() def test_csv_db_error_mid_result_returns_500(client): @@ -105,10 +111,7 @@ async def failing_rows(db, params): yield make_row() raise RuntimeError("connection lost") - with ( - patch.object(search, "iter_filtered_data", failing_rows), - patch.object(search, "has_more_rows_than", AsyncMock(return_value=False)), - ): + with patch.object(search, "iter_filtered_data", failing_rows): response = get_csv(client) assert response.status_code == 500 @@ -123,10 +126,44 @@ def test_csv_spills_large_output_to_disk(client, db_calls): assert all(len(dict(zip(HEADER, r))["protein_sequence"]) == 100_000 for r in rows[1:]) -async def test_cap_check_offset_stays_within_bigint(db): - from app.db.queries import has_more_rows_than +async def test_closing_iter_filtered_data_closes_the_cursor(db): + from contextlib import aclosing + + from app.db.queries import iter_filtered_data from app.models.query_params import CLEANSearchQueryParams - await has_more_rows_than(db, CLEANSearchQueryParams(offset=MAX_BIGINT - 1), search.MAX_SEARCH_LIMIT) + closed = [] + + async def iterate(query, *args): + try: + yield make_row() + yield make_row() + finally: + closed.append(True) + + db.iterate = iterate + async with aclosing(iter_filtered_data(db, CLEANSearchQueryParams(limit=2))) as rows: + await anext(rows) + + assert closed == [True] + + +async def test_write_csv_closes_rows_when_over_cap(): + from fastapi import HTTPException + + closed = [] + + async def rows(): + try: + for _ in range(5): + yield make_row() + finally: + closed.append(True) + + # Holding a reference means only an explicit close, not garbage collection, can close it + generator = rows() + with pytest.raises(HTTPException) as error: + await search.write_csv(generator, max_rows=3) - assert db.fetchval.call_args.args[0].rstrip().endswith(f"OFFSET {MAX_BIGINT})") + assert error.value.status_code == 400 + assert closed == [True] From eda98e6fe92f2262583299c8bffd02032519df79 Mon Sep 17 00:00:00 2001 From: matthewberry Date: Thu, 22 Oct 2026 12:18:52 -0500 Subject: [PATCH 5/5] feat: cap JSON search limit at the auto-pagination page size JSON is paginated, so an explicit JSON limit above AUTO_PAGINATION_THRESHOLD (5000) now gets a 422 shaped like the le= error on limit. CSV, the bulk download, keeps the 100,000 cap. A JSON request builds a Pydantic object per row and keeps the whole response in memory, so 100,000-row JSON requests were the most expensive /search requests. --- app/main.py | 3 ++- app/routers/search.py | 25 +++++++++++++++++++++---- tests/test_pagination_bounds.py | 3 ++- tests/test_search_csv.py | 9 +++++++++ 4 files changed, 34 insertions(+), 6 deletions(-) diff --git a/app/main.py b/app/main.py index 18ef53c..b6eca5e 100644 --- a/app/main.py +++ b/app/main.py @@ -53,7 +53,8 @@ async def lifespan(app: FastAPI): When a query would return more than {settings.AUTO_PAGINATION_THRESHOLD} records and no explicit limit is provided, the API will automatically paginate results to return {settings.AUTO_PAGINATION_THRESHOLD} records at a time. The response will include pagination -metadata with links to navigate to next and previous pages. +metadata with links to navigate to next and previous pages. An explicit JSON `limit` can be at +most {settings.AUTO_PAGINATION_THRESHOLD}. This threshold can be configured using the `AUTO_PAGINATION_THRESHOLD` environment variable. diff --git a/app/routers/search.py b/app/routers/search.py index c50d78d..9bde8a9 100644 --- a/app/routers/search.py +++ b/app/routers/search.py @@ -5,6 +5,7 @@ from urllib.parse import urlencode from fastapi import APIRouter, Depends, HTTPException, Query, Request +from fastapi.exceptions import RequestValidationError from fastapi.responses import StreamingResponse from starlette.background import BackgroundTask from loguru import logger @@ -117,7 +118,9 @@ def parse_query_params( default=ResponseFormat.JSON, description="Response format (json or csv)" ), limit: Optional[int] = Query( - None, ge=1, le=MAX_SEARCH_LIMIT, description="Maximum number of records to return" + None, ge=1, le=MAX_SEARCH_LIMIT, + description="Maximum number of records to return " + f"(at most {settings.AUTO_PAGINATION_THRESHOLD} for JSON, {MAX_SEARCH_LIMIT} for CSV)" ), offset: Optional[int] = Query(0, ge=0, le=MAX_BIGINT, description="Number of records to skip"), ordering: Optional[str] = Query( @@ -127,6 +130,19 @@ def parse_query_params( ), ) -> CLEANSearchQueryParams: """Parse and validate query parameters.""" + # JSON is paginated, so a JSON page is never larger than the default page size. + # Shaped like the le= error on limit, so both caps fail the same way. + max_json_limit = settings.AUTO_PAGINATION_THRESHOLD + if format == ResponseFormat.JSON and limit is not None and limit > max_json_limit: + raise RequestValidationError([{ + "type": "less_than_equal", + "loc": ("query", "limit"), + "msg": f"Input should be less than or equal to {max_json_limit} for JSON responses; " + "page with offset, or use format=csv", + "input": str(limit), + "ctx": {"le": max_json_limit}, + }]) + try: # Validate format if format not in [fmt for fmt in ResponseFormat]: @@ -171,9 +187,10 @@ async def get_data( with OR logic, while filters on different parameters are combined with AND logic. The response format can be either JSON (default) or CSV. JSON results are automatically -paginated when no explicit `limit` is provided. CSV results are not paginated: without a -`limit`, every matching record is returned, up to 100,000; a query that matches more is -rejected with a 400. +paginated when no explicit `limit` is provided, and an explicit JSON `limit` can be no larger +than that page size (`AUTO_PAGINATION_THRESHOLD`, 5,000 by default). CSV results are not +paginated: without a `limit`, every matching record is returned, up to 100,000; a query that +matches more is rejected with a 400. ### URL examples diff --git a/tests/test_pagination_bounds.py b/tests/test_pagination_bounds.py index f51d167..4c0d535 100644 --- a/tests/test_pagination_bounds.py +++ b/tests/test_pagination_bounds.py @@ -10,6 +10,7 @@ (SEARCH, "limit", -1), (SEARCH, "limit", 0), (SEARCH, "limit", 100001), + (SEARCH, "limit", 5001), # JSON cap; CSV allows up to 100000 (SEARCH, "offset", -5), (SEARCH, "offset", TOO_BIG), (TYPEAHEAD, "limit", -1), @@ -20,7 +21,7 @@ ] ACCEPTED = [ - (SEARCH, "limit", 100000), + (SEARCH, "limit", 5000), (SEARCH, "offset", TOO_BIG - 1), (TYPEAHEAD, "limit", 100), (TYPEAHEAD, "offset", TOO_BIG - 1), diff --git a/tests/test_search_csv.py b/tests/test_search_csv.py index a86f4f4..f7a1261 100644 --- a/tests/test_search_csv.py +++ b/tests/test_search_csv.py @@ -106,6 +106,15 @@ def test_csv_with_limit_is_not_capped(client, db_calls, monkeypatch): assert db_calls.streamed_params[0].offset == 100 +def test_csv_accepts_limits_above_json_cap(client, db_calls): + for limit in (settings.AUTO_PAGINATION_THRESHOLD + 1, search.MAX_SEARCH_LIMIT): + parse(get_csv(client, f"limit={limit}")) + + assert [p.limit for p in db_calls.streamed_params] == [ + settings.AUTO_PAGINATION_THRESHOLD + 1, search.MAX_SEARCH_LIMIT, + ] + + def test_csv_db_error_mid_result_returns_500(client): async def failing_rows(db, params): yield make_row()