Respect _size on arbitrary SQL query pages

datasette/templates/query.html

@@ -66,6 +66,9 @@ <h3>Query parameters</h3>
     <p>
         {% if not hide_sql %}<button id="sql-format" type="button" hidden>Format SQL</button>{% endif %}
         <input type="submit" value="Run SQL"{% if canned_query_write and db_is_immutable %} disabled{% endif %}>
+        {% if query.params.get("_size") %}
+            <input type="hidden" name="_size" value="{{ query.params.get("_size") }}">
+        {% endif %}
         {{ show_hide_hidden }}
         {% if canned_query and edit_sql_url %}<a href="{{ edit_sql_url }}" class="canned-query-edit-sql">Edit SQL</a>{% endif %}
     </p>

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,

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(

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