From fd368d3b2c5a5d9c3e10a21638f6ea9a71471b52 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 1 Jun 2021 09:12:32 -0700
Subject: [PATCH] New _nocount=1 option, used to speed up CSVs - closes #1353

---
 datasette/views/base.py  | 15 +++++++++++----
 datasette/views/table.py |  6 +++++-
 docs/json_api.rst        |  3 +++
 tests/test_api.py        |  9 +++++++++
 tests/test_csv.py        |  6 ++++++
 5 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/datasette/views/base.py b/datasette/views/base.py
index b8c581fc3a..26edfde553 100644
--- a/datasette/views/base.py
+++ b/datasette/views/base.py
@@ -263,12 +263,19 @@ async def get(self, request, db_name, **kwargs):
 
     async def as_csv(self, request, database, hash, **kwargs):
         stream = request.args.get("_stream")
-        # Do not calculate facets:
-        if not request.args.get("_nofacet"):
+        # Do not calculate facets or counts:
+        extra_parameters = [
+            "{}=1".format(key)
+            for key in ("_nofacet", "_nocount")
+            if not request.args.get(key)
+        ]
+        if extra_parameters:
             if not request.query_string:
-                new_query_string = "_nofacet=1"
+                new_query_string = "&".join(extra_parameters)
             else:
-                new_query_string = request.query_string + "&_nofacet=1"
+                new_query_string = (
+                    request.query_string + "&" + "&".join(extra_parameters)
+                )
             new_scope = dict(
                 request.scope, query_string=new_query_string.encode("latin-1")
             )
diff --git a/datasette/views/table.py b/datasette/views/table.py
index 7fbf670b43..d47865f0e5 100644
--- a/datasette/views/table.py
+++ b/datasette/views/table.py
@@ -697,7 +697,11 @@ async def data(
             except KeyError:
                 pass
 
-        if count_sql and filtered_table_rows_count is None:
+        if (
+            count_sql
+            and filtered_table_rows_count is None
+            and not request.args.get("_nocount")
+        ):
             try:
                 count_rows = list(await db.execute(count_sql, from_sql_params))
                 filtered_table_rows_count = count_rows[0][0]
diff --git a/docs/json_api.rst b/docs/json_api.rst
index f1c347b753..660fbc1ca9 100644
--- a/docs/json_api.rst
+++ b/docs/json_api.rst
@@ -386,6 +386,9 @@ Special table arguments
 ``?_nofacet=1``
     Disable all facets and facet suggestions for this page, including any defined by :ref:`facets_metadata`.
 
+``?_nocount=1``
+    Disable the ``select count(*)`` query used on this page - a count of ``None`` will be returned instead.
+
 ``?_trace=1``
     Turns on tracing for this page: SQL queries executed during the request will
     be gathered and included in the response, either in a new ``"_traces"`` key
diff --git a/tests/test_api.py b/tests/test_api.py
index 5e63913341..49b3bbe992 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -1683,6 +1683,15 @@ def test_nofacet(app_client, nofacet):
         assert response.json["facet_results"] != {}
 
 
+@pytest.mark.parametrize("nocount,expected_count", ((True, None), (False, 15)))
+def test_nocount(app_client, nocount, expected_count):
+    path = "/fixtures/facetable.json"
+    if nocount:
+        path += "?_nocount=1"
+    response = app_client.get(path)
+    assert response.json["filtered_table_rows_count"] == expected_count
+
+
 def test_expand_labels(app_client):
     response = app_client.get(
         "/fixtures/facetable.json?_shape=object&_labels=1&_size=2"
diff --git a/tests/test_csv.py b/tests/test_csv.py
index 40549fd857..02fe576675 100644
--- a/tests/test_csv.py
+++ b/tests/test_csv.py
@@ -175,3 +175,9 @@ def test_table_csv_stream_does_not_calculate_facets(app_client):
     response = app_client.get("/fixtures/simple_primary_key.csv?_trace=1")
     soup = Soup(response.text, "html.parser")
     assert "select content, count(*) as n" not in soup.find("pre").text
+
+
+def test_table_csv_stream_does_not_calculate_counts(app_client):
+    response = app_client.get("/fixtures/simple_primary_key.csv?_trace=1")
+    soup = Soup(response.text, "html.parser")
+    assert "select count(*)" not in soup.find("pre").text