From 9a5c92b63e7917c93cc502478493c51c781b2ecc Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 2 Jan 2021 14:03:52 -0800
Subject: [PATCH] db.enable_counts() method, closes #213

---
 docs/python-api.rst         |  6 ++++++
 sqlite_utils/db.py          | 23 +++++++++++++++++++++--
 tests/test_enable_counts.py | 29 ++++++++++++++++++++++++++++-
 3 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index 400fa58..1fc5ba8 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -1706,6 +1706,12 @@ This will create the ``_counts`` table if it does not already exist, with the fo
 
 Once enabled, table counts can be accessed by querying the ``counts`` table. The count records will be automatically kept up-to-date by the triggers when rows are added or deleted to the table.
 
+You can enable cached counts for every table in a database (except for virtual tables and the ``_counts`` table itself) using the database ``enable_counts()`` method:
+
+.. code-block:: python
+
+    db.enable_counts()
+
 Creating indexes
 ================
 
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index e2f55bf..3118138 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -145,6 +145,11 @@ class InvalidColumns(Exception):
 
 
 class Database:
+    _counts_table_name = "_counts"
+    _counts_table_create = "CREATE TABLE IF NOT EXISTS [{}]([table] TEXT PRIMARY KEY, count INTEGER DEFAULT 0);".format(
+        _counts_table_name
+    )
+
     def __init__(
         self,
         filename_or_conn=None,
@@ -279,6 +284,19 @@ class Database:
         if self.journal_mode != "delete":
             self.execute("PRAGMA journal_mode=delete;")
 
+    def _ensure_counts_table(self):
+        with self.conn:
+            self.execute(self._counts_table_create)
+
+    def enable_counts(self):
+        self._ensure_counts_table()
+        for table in self.tables:
+            if (
+                table.virtual_table_using is None
+                and table.name != self._counts_table_name
+            ):
+                table.enable_counts()
+
     def execute_returning_dicts(self, sql, params=None):
         cursor = self.execute(sql, params or tuple())
         keys = [d[0] for d in cursor.description]
@@ -1178,7 +1196,7 @@ class Table(Queryable):
         sql = (
             textwrap.dedent(
                 """
-        CREATE TABLE IF NOT EXISTS [{counts_table}]([table] TEXT PRIMARY KEY, count INTEGER DEFAULT 0);
+        {create_counts_table}
         CREATE TRIGGER IF NOT EXISTS [{table}{counts_table}_insert] AFTER INSERT ON [{table}]
         BEGIN
             INSERT OR REPLACE INTO [{counts_table}]
@@ -1206,7 +1224,8 @@ class Table(Queryable):
             )
             .strip()
             .format(
-                counts_table="_counts",
+                create_counts_table=self.db._counts_table_create,
+                counts_table=self.db._counts_table_name,
                 table=self.name,
                 table_quoted=self.db.escape(self.name),
             )
diff --git a/tests/test_enable_counts.py b/tests/test_enable_counts.py
index 9dc2d54..d876a35 100644
--- a/tests/test_enable_counts.py
+++ b/tests/test_enable_counts.py
@@ -1,4 +1,4 @@
-def test_enable_counts(fresh_db):
+def test_enable_counts_specific_table(fresh_db):
     foo = fresh_db["foo"]
     assert fresh_db.table_names() == []
     for i in range(10):
@@ -25,3 +25,30 @@ def test_enable_counts(fresh_db):
     foo.delete_where()
     assert foo.count == 0
     assert list(fresh_db["_counts"].rows) == [{"count": 0, "table": "foo"}]
+
+
+def test_enable_counts_all_tables(fresh_db):
+    foo = fresh_db["foo"]
+    bar = fresh_db["bar"]
+    foo.insert({"name": "Cleo"})
+    bar.insert({"name": "Cleo"})
+    foo.enable_fts(["name"])
+    fresh_db.enable_counts()
+    assert set(fresh_db.table_names()) == {
+        "foo",
+        "bar",
+        "foo_fts",
+        "foo_fts_data",
+        "foo_fts_idx",
+        "foo_fts_docsize",
+        "foo_fts_config",
+        "_counts",
+    }
+    assert list(fresh_db["_counts"].rows) == [
+        {"count": 1, "table": "foo"},
+        {"count": 1, "table": "bar"},
+        {"count": 3, "table": "foo_fts_data"},
+        {"count": 1, "table": "foo_fts_idx"},
+        {"count": 1, "table": "foo_fts_docsize"},
+        {"count": 1, "table": "foo_fts_config"},
+    ]