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

@@ -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"

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)

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

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(