table.enable_counts() method, closes #212

docs/python-api.rst

@@ -1682,6 +1682,30 @@ This runs the following SQL::
 
     INSERT INTO dogs_fts (dogs_fts) VALUES ("optimize");
 
+.. _python_api_enable_counts:
+
+Enabling cached counts for a table
+==================================
+
+The ``select count(*)`` query in SQLite requires a full scan of the primary key index, and can take an increasingly long time as the table grows larger.
+
+The ``table.enable_counts()`` method can be used to configure triggers to continuously update a record in a ``_counts`` table. This value can then be used to quickly retrieve the count of rows in the associated table.
+
+.. code-block:: python
+
+    db["dogs"].enable_counts()
+
+This will create the ``_counts`` table if it does not already exist, with the following schema:
+
+.. code-block:: sql
+
+    CREATE TABLE [_counts] (
+        [table] TEXT PRIMARY KEY,
+        [count] INTEGER DEFAULT 0
+    )
+
+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.
+
 Creating indexes
 ================
 

sqlite_utils/db.py

@@ -1174,6 +1174,46 @@ def add_foreign_key(
         self.db.add_foreign_keys([(self.name, column, other_table, other_column)])
         return self
 
+    def enable_counts(self):
+        sql = (
+            textwrap.dedent(
+                """
+        CREATE TABLE IF NOT EXISTS [{counts_table}]([table] TEXT PRIMARY KEY, count INTEGER DEFAULT 0);
+        CREATE TRIGGER IF NOT EXISTS [{table}{counts_table}_insert] AFTER INSERT ON [{table}]
+        BEGIN
+            INSERT OR REPLACE INTO [{counts_table}]
+            VALUES (
+                {table_quoted},
+                COALESCE(
+                    (SELECT count FROM [{counts_table}] WHERE [table] = {table_quoted}),
+                0
+                ) + 1
+            );
+        END;
+        CREATE TRIGGER IF NOT EXISTS [{table}{counts_table}_delete] AFTER DELETE ON [{table}]
+        BEGIN
+            INSERT OR REPLACE INTO [{counts_table}]
+            VALUES (
+                {table_quoted},
+                COALESCE(
+                    (SELECT count FROM [{counts_table}] WHERE [table] = {table_quoted}),
+                0
+                ) - 1
+            );
+        END;
+        INSERT OR REPLACE INTO _counts VALUES ({table_quoted}, (select count(*) from [{table}]));
+        """
+            )
+            .strip()
+            .format(
+                counts_table="_counts",
+                table=self.name,
+                table_quoted=self.db.escape(self.name),
+            )
+        )
+        with self.db.conn:
+            self.db.conn.executescript(sql)
+
     def enable_fts(
         self,
         columns,

tests/test_enable_counts.py

@@ -0,0 +1,27 @@
+def test_enable_counts(fresh_db):
+    foo = fresh_db["foo"]
+    assert fresh_db.table_names() == []
+    for i in range(10):
+        foo.insert({"name": "item {}".format(i)})
+    assert fresh_db.table_names() == ["foo"]
+    assert foo.count == 10
+    # Now enable counts
+    foo.enable_counts()
+    assert foo.triggers_dict == {
+        "foo_counts_insert": "CREATE TRIGGER [foo_counts_insert] AFTER INSERT ON [foo]\nBEGIN\n    INSERT OR REPLACE INTO [_counts]\n    VALUES (\n        'foo',\n        COALESCE(\n            (SELECT count FROM [_counts] WHERE [table] = 'foo'),\n        0\n        ) + 1\n    );\nEND",
+        "foo_counts_delete": "CREATE TRIGGER [foo_counts_delete] AFTER DELETE ON [foo]\nBEGIN\n    INSERT OR REPLACE INTO [_counts]\n    VALUES (\n        'foo',\n        COALESCE(\n            (SELECT count FROM [_counts] WHERE [table] = 'foo'),\n        0\n        ) - 1\n    );\nEND",
+    }
+    assert fresh_db.table_names() == ["foo", "_counts"]
+    assert list(fresh_db["_counts"].rows) == [{"count": 10, "table": "foo"}]
+    # Add some items to test the triggers
+    for i in range(5):
+        foo.insert({"name": "item {}".format(10 + i)})
+    assert foo.count == 15
+    assert list(fresh_db["_counts"].rows) == [{"count": 15, "table": "foo"}]
+    # Delete some items
+    foo.delete_where("rowid < 7")
+    assert foo.count == 9
+    assert list(fresh_db["_counts"].rows) == [{"count": 9, "table": "foo"}]
+    foo.delete_where()
+    assert foo.count == 0
+    assert list(fresh_db["_counts"].rows) == [{"count": 0, "table": "foo"}]