From 9ecce07b083824f56bd96966a1f63e18d44489b1 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 21 Aug 2024 19:09:25 -0700
Subject: [PATCH] count all rows button on table page, refs #2408

---
 datasette/templates/table.html | 33 ++++++++++++++++++++++++++++++++-
 datasette/url_builder.py       |  6 ++++++
 datasette/views/table.py       | 11 ++++++++++-
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/datasette/templates/table.html b/datasette/templates/table.html
index 35e0b9c1..187f0143 100644
--- a/datasette/templates/table.html
+++ b/datasette/templates/table.html
@@ -40,7 +40,10 @@
 {% endif %}
 
 {% if count or human_description_en %}
-    <h3>{% if count or count == 0 %}{{ "{:,}".format(count) }} row{% if count == 1 %}{% else %}s{% endif %}{% endif %}
+    <h3>
+        {% if count == count_limit + 1 %}&gt;{{ "{:,}".format(count_limit) }} rows
+        {% if allow_execute_sql and query.sql %} <a class="count-sql" style="font-size: 0.8em; padding-left: 0.5em" href="{{ urls.database_query(database, count_sql) }}">count all rows</a>{% endif %}
+        {% elif count or count == 0 %}{{ "{:,}".format(count) }} row{% if count == 1 %}{% else %}s{% endif %}{% endif %}
         {% if human_description_en %}{{ human_description_en }}{% endif %}
     </h3>
 {% endif %}
@@ -172,4 +175,32 @@
     <pre class="wrapped-sql">{{ view_definition }}</pre>
 {% endif %}
 
+{% if allow_execute_sql and query.sql %}
+<script>
+document.addEventListener('DOMContentLoaded', function() {
+    const countLink = document.querySelector('a.count-sql');
+    if (countLink) {
+        countLink.addEventListener('click', function(ev) {
+            ev.preventDefault();
+            // Replace countLink with span with same style attribute
+            const span = document.createElement('span');
+            span.textContent = 'counting...';
+            span.setAttribute('style', countLink.getAttribute('style'));
+            countLink.replaceWith(span);
+            countLink.setAttribute('disabled', 'disabled');
+            let url = countLink.href.replace(/(\?|$)/, '.json$1');
+            fetch(url)
+                .then(response => response.json())
+                .then(data => {
+                    const count = data['rows'][0]['count(*)'];
+                    const formattedCount = count.toLocaleString();
+                    span.closest('h3').textContent = formattedCount + ' rows';
+                })
+                .catch(error => countLink.textContent = 'error');
+        });
+    }
+});
+</script>
+{% endif %}
+
 {% endblock %}
diff --git a/datasette/url_builder.py b/datasette/url_builder.py
index 9c6bbde0..16b3d42b 100644
--- a/datasette/url_builder.py
+++ b/datasette/url_builder.py
@@ -31,6 +31,12 @@ class Urls:
         db = self.ds.get_database(database)
         return self.path(tilde_encode(db.route), format=format)
 
+    def database_query(self, database, sql, format=None):
+        path = f"{self.database(database)}/-/query?" + urllib.parse.urlencode(
+            {"sql": sql}
+        )
+        return self.path(path, format=format)
+
     def table(self, database, table, format=None):
         path = f"{self.database(database)}/{tilde_encode(table)}"
         if format is not None:
diff --git a/datasette/views/table.py b/datasette/views/table.py
index d71efeb0..ea044b36 100644
--- a/datasette/views/table.py
+++ b/datasette/views/table.py
@@ -929,6 +929,7 @@ async def table_view_traced(datasette, request):
                         database=resolved.db.name,
                         table=resolved.table,
                     ),
+                    count_limit=resolved.db.count_limit,
                 ),
                 request=request,
                 view_name="table",
@@ -1280,6 +1281,9 @@ async def table_view_data(
     if extra_extras:
         extras.update(extra_extras)
 
+    async def extra_count_sql():
+        return count_sql
+
     async def extra_count():
         "Total count of rows matching these filters"
         # Calculate the total count for this query
@@ -1299,8 +1303,11 @@ async def table_view_data(
 
         # Otherwise run a select count(*) ...
         if count_sql and count is None and not nocount:
+            count_sql_limited = (
+                f"select count(*) from (select * {from_sql} limit 10001)"
+            )
             try:
-                count_rows = list(await db.execute(count_sql, from_sql_params))
+                count_rows = list(await db.execute(count_sql_limited, from_sql_params))
                 count = count_rows[0][0]
             except QueryInterrupted:
                 pass
@@ -1615,6 +1622,7 @@ async def table_view_data(
             "facet_results",
             "facets_timed_out",
             "count",
+            "count_sql",
             "human_description_en",
             "next_url",
             "metadata",
@@ -1647,6 +1655,7 @@ async def table_view_data(
 
     registry = Registry(
         extra_count,
+        extra_count_sql,
         extra_facet_results,
         extra_facets_timed_out,
         extra_suggested_facets,