Merge 28c411679c into c435a80030

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(