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(