Return full CSV search results up to 100,000 rows; cap JSON limit at 5,000

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:

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:

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

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)}")

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),