diff --git a/app/db/database.py b/app/db/database.py index 85ea9d3..80a6a49 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 @@ -54,6 +54,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 879b024..8c2ff25 100644 --- a/app/db/queries.py +++ b/app/db/queries.py @@ -1,4 +1,5 @@ -from typing import Any, Dict, List, Tuple +from contextlib import aclosing +from typing import Any, AsyncIterator, Dict, List, Tuple import re from app.db.database import Database @@ -136,10 +137,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 = """ @@ -170,11 +169,26 @@ async def get_filtered_data( query += f" OFFSET {params.offset}" # 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) + # 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/main.py b/app/main.py index cf72447..b6eca5e 100644 --- a/app/main.py +++ b/app/main.py @@ -53,10 +53,15 @@ 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. +CSV responses are not paginated. Without an explicit limit, a CSV search returns every matching +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 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 bcaa61a..9bde8a9 100644 --- a/app/routers/search.py +++ b/app/routers/search.py @@ -1,15 +1,18 @@ import csv -from io import StringIO -from typing import Any, List, Literal, Optional +from contextlib import aclosing +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.exceptions import RequestValidationError from fastapi.responses import StreamingResponse +from starlette.background import BackgroundTask from loguru import logger 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, 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 @@ -18,6 +21,57 @@ # 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 +CSV_SPOOL_MAX_MEMORY = 8 * 1024 * 1024 + + +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="") + 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 + + +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( # String filters @@ -64,7 +118,9 @@ 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 " + 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( @@ -74,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]: @@ -117,8 +186,11 @@ 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, 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 @@ -207,41 +279,33 @@ 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() - - # 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) + if params.limit is None: + # 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( - iter([output.getvalue()]), + 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 + 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, @@ -314,6 +378,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_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 12d9cad..f7a1261 100644 --- a/tests/test_search_csv.py +++ b/tests/test_search_csv.py @@ -2,18 +2,45 @@ from io import StringIO from unittest.mock import AsyncMock, patch +import pytest + import app.routers.search as search +from app.core.config import settings from app.models.clean_data import CLEANColumn HEADER = [c.value for c in CLEANColumn] -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 = [] + rows_closed = False + + async def iter_filtered_data(db, params): + Calls.streamed_params.append(params.model_copy()) + try: + for row in Calls.rows: + yield row + finally: + Calls.rows_closed = True + 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, "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.get_total_count, Calls.get_filtered_data = 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))) @@ -25,12 +52,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])) @@ -40,5 +68,111 @@ 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, 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 [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.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, 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 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_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 + + +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() + raise RuntimeError("connection lost") + + with patch.object(search, "iter_filtered_data", failing_rows): + 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:]) + + +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 + + 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 error.value.status_code == 400 + assert closed == [True]