From 28c411679c8ca919dfb0fb491e8f5f13eb0ff2b9 Mon Sep 17 00:00:00 2001 From: kokokoXUY <13682395396@163.com> Date: Sat, 26 Sep 2026 11:41:45 +0800 Subject: [PATCH] fix: report custom query CSV truncation in response headers Expose the row limit and truncation flag so clients can detect incomplete CSV exports. Keep custom query CSV within the configured limit when it equals the page size. Fixes #1553. --- datasette/views/base.py | 6 ++++ datasette/views/database.py | 15 ++++++-- docs/csv_export.rst | 7 ++++ tests/test_csv.py | 69 +++++++++++++++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 2 deletions(-) diff --git a/datasette/views/base.py b/datasette/views/base.py index 7262edb4..ce74c039 100644 --- a/datasette/views/base.py +++ b/datasette/views/base.py @@ -327,6 +327,12 @@ async def stream_csv(datasette, fetch_data, request, database): headers = {} if datasette.cors: add_cors_headers(headers) + if "truncated" in data: + headers["Datasette-Truncated"] = str(data["truncated"]).lower() + headers["Datasette-Max-Returned-Rows"] = str(datasette.max_returned_rows) + if datasette.cors: + headers["Access-Control-Expose-Headers"] += ", Datasette-Truncated" + headers["Access-Control-Expose-Headers"] += ", Datasette-Max-Returned-Rows" if request.args.get("_dl", None): if not trace: content_type = "text/csv; charset=utf-8" diff --git a/datasette/views/database.py b/datasette/views/database.py index da207601..8b8e296b 100644 --- a/datasette/views/database.py +++ b/datasette/views/database.py @@ -862,8 +862,19 @@ class QueryView(View): async def fetch_data_for_csv(request, _next=None): # Reuse the trusted magic parameter values prepared above. - results = await db.execute(sql, params_for_query, truncate=True) - data = {"rows": results.rows, "columns": results.columns} + # CSV queries do not paginate, so the page-size sentinel must not + # increase the configured maximum by one. + results = await db.execute( + sql, + params_for_query, + truncate=True, + page_size=datasette.max_returned_rows + 1, + ) + data = { + "rows": results.rows, + "columns": results.columns, + "truncated": results.truncated, + } return data, None, None return await stream_csv(datasette, fetch_data_for_csv, request, db.name) diff --git a/docs/csv_export.rst b/docs/csv_export.rst index 34d81594..82311465 100644 --- a/docs/csv_export.rst +++ b/docs/csv_export.rst @@ -45,6 +45,13 @@ The following options can be used to customize the CSVs returned by Datasette. ``?_dl=on`` Causes Datasette to return a ``content-disposition: attachment; filename="filename.csv"`` header. +Custom SQL query CSV responses include a ``Datasette-Truncated`` response +header with the value ``true`` or ``false``. A value of ``true`` means rows +were omitted because of the :ref:`setting_max_returned_rows` limit. The +``Datasette-Max-Returned-Rows`` header reports that limit, including ``0`` +when it is disabled. These headers are also available to cross-origin browser +clients when CORS is enabled. + Streaming all records --------------------- diff --git a/tests/test_csv.py b/tests/test_csv.py index adae7e24..10aed9c8 100644 --- a/tests/test_csv.py +++ b/tests/test_csv.py @@ -49,6 +49,7 @@ async def test_table_csv(ds_client): response = await ds_client.get("/fixtures/simple_primary_key.csv?_oh=1") assert response.status_code == 200 assert not response.headers.get("Access-Control-Allow-Origin") + assert "Datasette-Truncated" not in response.headers assert response.headers["content-type"] == "text/plain; charset=utf-8" assert response.text == EXPECTED_TABLE_CSV @@ -166,6 +167,74 @@ async def test_custom_sql_csv(ds_client): assert response.text == EXPECTED_CUSTOM_CSV +@pytest.mark.asyncio +@pytest.mark.parametrize("download", (False, True)) +@pytest.mark.parametrize( + "sql,truncated", + [ + ("select content from simple_primary_key limit 2", "false"), + ("select * from no_primary_key", "true"), + ], +) +async def test_custom_sql_csv_reports_row_limit(ds_client, sql, truncated, download): + parameters = {"sql": sql} + if download: + parameters["_dl"] = "1" + response = await ds_client.get( + "/fixtures/-/query.csv?" + urllib.parse.urlencode(parameters) + ) + assert response.status_code == 200 + assert response.headers["Datasette-Truncated"] == truncated + assert response.headers["Datasette-Max-Returned-Rows"] == "100" + + +@pytest.mark.asyncio +async def test_custom_sql_csv_row_limit_headers_are_exposed_to_cors(): + ds = Datasette(memory=True, cors=True) + try: + response = await ds.client.get("/_memory/-/query.csv?sql=select+1") + assert response.status_code == 200 + assert response.headers["Datasette-Truncated"] == "false" + assert response.headers["Access-Control-Expose-Headers"] == ( + "Link, Datasette-Truncated, Datasette-Max-Returned-Rows" + ) + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_custom_sql_csv_without_row_limit(): + ds = Datasette(memory=True, settings={"max_returned_rows": 0}) + try: + response = await ds.client.get( + "/_memory/-/query.csv?sql=select+1+as+n+union+all+select+2" + ) + assert response.status_code == 200 + assert response.text == "n\r\n1\r\n2\r\n" + assert response.headers["Datasette-Truncated"] == "false" + assert response.headers["Datasette-Max-Returned-Rows"] == "0" + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_custom_sql_csv_row_limit_matching_page_size(): + ds = Datasette( + memory=True, settings={"max_returned_rows": 2, "default_page_size": 2} + ) + try: + response = await ds.client.get( + "/_memory/-/query.csv?sql=select+1+as+n+union+all+select+2+" + "union+all+select+3+union+all+select+4" + ) + assert response.status_code == 200 + assert response.text == "n\r\n1\r\n2\r\n" + assert response.headers["Datasette-Truncated"] == "true" + assert response.headers["Datasette-Max-Returned-Rows"] == "2" + finally: + ds.close() + + @pytest.mark.asyncio @pytest.mark.parametrize("download", (False, True)) @pytest.mark.parametrize(