{% if not hide_sql %}{% endif %} + {% if query.params.get("_size") %} + + {% endif %} {{ show_hide_hidden }} {% if canned_query and edit_sql_url %}Edit SQL{% endif %}
diff --git a/datasette/views/database.py b/datasette/views/database.py index faf870d032..731719d8e1 100644 --- a/datasette/views/database.py +++ b/datasette/views/database.py @@ -31,7 +31,7 @@ truncate_url, InvalidSql, ) -from datasette.utils.asgi import AsgiFileDownload, NotFound, Response, Forbidden +from datasette.utils.asgi import AsgiFileDownload, BadRequest, NotFound, Response, Forbidden from datasette.plugins import pm from .base import BaseView, DatasetteError, View, _error, stream_csv @@ -311,6 +311,9 @@ class QueryContext(Context): columns: list = field(metadata={"help": "List of column names"}) renderers: dict = field(metadata={"help": "Dictionary of renderer name to URL"}) url_csv: str = field(metadata={"help": "URL for CSV export"}) + truncated: bool = field( + metadata={"help": "Boolean indicating if the query results were truncated"} + ) show_hide_hidden: str = field( metadata={"help": "Hidden input field for the _show_sql parameter"} ) @@ -590,11 +593,25 @@ async def get(self, request, datasette): extra_args = {} if params.get("_timelimit"): extra_args["custom_time_limit"] = int(params["_timelimit"]) + page_size = None + if params.get("_size"): + page_size = params["_size"] + if page_size == "max": + page_size = datasette.max_returned_rows + try: + page_size = int(page_size) + if page_size < 0: + raise ValueError + except ValueError: + raise BadRequest("_size must be a positive integer") + if page_size > datasette.max_returned_rows: + raise BadRequest(f"_size must be <= {datasette.max_returned_rows}") format_ = request.url_vars.get("format") or "html" query_error = None results = None + truncated = False rows = [] columns = [] @@ -614,6 +631,10 @@ async def get(self, request, datasette): ) columns = results.columns rows = results.rows + truncated = results.truncated + if page_size is not None and len(rows) > page_size: + rows = rows[:page_size] + truncated = True except QueryInterrupted as ex: raise DatasetteError( textwrap.dedent(""" @@ -647,7 +668,10 @@ async def get(self, request, datasette): async def fetch_data_for_csv(request, _next=None): results = await db.execute(sql, params, truncate=True) - data = {"rows": results.rows, "columns": results.columns} + csv_rows = results.rows + if page_size is not None and len(csv_rows) > page_size: + csv_rows = csv_rows[:page_size] + data = {"rows": csv_rows, "columns": results.columns} return data, None, None return await stream_csv(datasette, fetch_data_for_csv, request, db.name) @@ -665,7 +689,7 @@ async def fetch_data_for_csv(request, _next=None): table=None, request=request, view_name="table", - truncated=results.truncated if results else False, + truncated=truncated, error=query_error, # These will be deprecated in Datasette 1.0: args=request.args, @@ -843,6 +867,7 @@ async def query_actions(): request=request, format="csv", extra_qs={"_size": "max"} ) ), + truncated=truncated, show_hide_hidden=markupsafe.Markup(show_hide_hidden), metadata=canned_query or metadata, alternate_url_json=alternate_url_json, diff --git a/tests/test_api.py b/tests/test_api.py index 392030d64b..c4832008d3 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -317,6 +317,34 @@ async def test_custom_sql(ds_client): } +@pytest.mark.asyncio +async def test_custom_sql_size(ds_client): + response = await ds_client.get( + "/fixtures/-/query.json?sql=select+*+from+compound_three_primary_keys&_size=10", + ) + data = response.json() + assert response.status_code == 200 + assert len(data["rows"]) == 10 + assert data["truncated"] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "size,expected_error", + [ + ("dog", "_size must be a positive integer"), + ("-1", "_size must be a positive integer"), + ("101", "_size must be <= 100"), + ], +) +async def test_custom_sql_size_validation(ds_client, size, expected_error): + response = await ds_client.get( + f"/fixtures/-/query.json?sql=select+1&_size={size}", + ) + assert response.status_code == 400 + assert response.json()["error"] == expected_error + + @pytest.mark.xfail(reason="Sometimes flaky in CI due to timing issues") def test_sql_time_limit(app_client_shorter_time_limit): response = app_client_shorter_time_limit.get( diff --git a/tests/test_html.py b/tests/test_html.py index 7425692d82..4209de5dc7 100644 --- a/tests/test_html.py +++ b/tests/test_html.py @@ -323,6 +323,19 @@ async def test_query_json_csv_export_links(ds_client): ) +@pytest.mark.asyncio +async def test_query_size(ds_client): + response = await ds_client.get( + "/fixtures/-/query?sql=select+*+from+compound_three_primary_keys&_size=10" + ) + assert response.status_code == 200 + soup = Soup(response.text, "html.parser") + assert "Custom SQL query returning more than 10 rows" in soup.find("h3").text + assert len(soup.select("table.rows-and-columns tbody tr")) == 10 + hidden = soup.select_one('input[type="hidden"][name="_size"]') + assert hidden["value"] == "10" + + @pytest.mark.asyncio async def test_query_parameter_form_fields(ds_client): response = await ds_client.get("/fixtures/-/query?sql=select+:name")