table.count now uses cached counts if db.use_counts_table

Closes #215

docs/cli.rst

@@ -1062,7 +1062,7 @@ Use the ``--sql`` option to output the SQL that would be executed, rather than r
 Enabling cached counts
 ======================
 
-``select count(*)`` queries can take a long time against large tables. ``sqlite-utils`` can speed these up by adding triggers to maintain a ``_counts`` table, see :ref:`python_api_enable_counts`.
+``select count(*)`` queries can take a long time against large tables. ``sqlite-utils`` can speed these up by adding triggers to maintain a ``_counts`` table, see :ref:`python_api_cached_table_counts` for details.
 
 The ``sqlite-utils enable-counts`` command can be used to configure these triggers, either for every table in the database or for specific tables.
 

docs/python-api.rst

@@ -1386,6 +1386,8 @@ The ``.count`` property shows the current number of rows (``select count(*) from
     >>> db["Street_Tree_List"].count
     189144
 
+This property will take advantage of :ref:`python_api_cached_table_counts` if the ``use_counts_table`` property is set on the database. You can avoid that optimization entirely by calling ``table.execute_count()`` instead of accessing the property.
+
 The ``.columns`` property shows the columns in the table or view::
 
     >>> db["PlantType"].columns
@@ -1700,9 +1702,9 @@ This runs the following SQL::
 
     INSERT INTO dogs_fts (dogs_fts) VALUES ("optimize");
 
-.. _python_api_enable_counts:
+.. _python_api_cached_table_counts:
 
-Enabling cached counts for a table
+Cached table counts using triggers
 ==================================
 
 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.
@@ -1718,18 +1720,44 @@ This will create the ``_counts`` table if it does not already exist, with the fo
 .. code-block:: sql
 
     CREATE TABLE [_counts] (
-        [table] TEXT PRIMARY KEY,
-        [count] INTEGER DEFAULT 0
+       [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.
-
 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()
 
+Once enabled, table counts will be stored in 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.
+
+To access these counts you can query the ``_counts`` table directly or you can use the ``db.cached_counts()`` method. This method returns a dictionary mapping tables to their counts::
+
+    >>> db.cached_counts()
+    {'global-power-plants': 33643,
+     'global-power-plants_fts_data': 136,
+     'global-power-plants_fts_idx': 199,
+     'global-power-plants_fts_docsize': 33643,
+     'global-power-plants_fts_config': 1}
+
+You can pass a list of table names to this method to retrieve just those counts::
+
+    >>> db.cached_counts(["global-power-plants"])
+    {'global-power-plants': 33643}
+
+The ``table.count`` property executes a ``select count(*)`` query by default, unless the ``db.use_counts_table`` property is set to ``True``.
+
+You can set ``use_counts_table`` to ``True`` when you instantiate the database object:
+
+.. code-block:: python
+
+    db = Database("global-power-plants.db", use_counts_table=True)
+
+If the property is ``True`` any calls to the ``table.count`` property will first attempt to find the cached count in the ``_counts`` table, and fall back on a ``count(*)`` query if the value is not available or the table is missing.
+
+Calling the ``.enable_counts()`` method on a database or table object will set ``use_counts_table`` to ``True`` for the lifetime of that database object.
+
 Creating indexes
 ================
 

sqlite_utils/db.py

@@ -144,11 +144,17 @@ class InvalidColumns(Exception):
     pass
 
 
+_COUNTS_TABLE_CREATE_SQL = """
+CREATE TABLE IF NOT EXISTS [{}](
+   [table] TEXT PRIMARY KEY,
+   count INTEGER DEFAULT 0
+);
+""".strip()
+
+
 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
-    )
+    use_counts_table = False
 
     def __init__(
         self,
@@ -157,6 +163,7 @@ class Database:
         recreate=False,
         recursive_triggers=True,
         tracer=None,
+        use_counts_table=False,
     ):
         assert (filename_or_conn is not None and not memory) or (
             filename_or_conn is None and memory
@@ -174,6 +181,7 @@ class Database:
         if recursive_triggers:
             self.execute("PRAGMA recursive_triggers=on;")
         self._registered_functions = set()
+        self.use_counts_table = use_counts_table
 
     @contextlib.contextmanager
     def tracer(self, tracer=None):
@@ -291,7 +299,7 @@ class Database:
 
     def _ensure_counts_table(self):
         with self.conn:
-            self.execute(self._counts_table_create)
+            self.execute(_COUNTS_TABLE_CREATE_SQL.format(self._counts_table_name))
 
     def enable_counts(self):
         self._ensure_counts_table()
@@ -301,6 +309,16 @@ class Database:
                 and table.name != self._counts_table_name
             ):
                 table.enable_counts()
+        self.use_counts_table = True
+
+    def cached_counts(self, tables=None):
+        sql = "select [table], count from {}".format(self._counts_table_name)
+        if tables:
+            sql += " where [table] in ({})".format(", ".join("?" for table in tables))
+        try:
+            return {r[0]: r[1] for r in self.execute(sql, tables).fetchall()}
+        except OperationalError:
+            return {}
 
     def execute_returning_dicts(self, sql, params=None):
         cursor = self.execute(sql, params or tuple())
@@ -602,12 +620,15 @@ class Queryable:
         self.db = db
         self.name = name
 
-    @property
-    def count(self):
+    def execute_count(self):
         return self.db.execute(
             "select count(*) from [{}]".format(self.name)
         ).fetchone()[0]
 
+    @property
+    def count(self):
+        return self.execute_count()
+
     @property
     def rows(self):
         return self.rows_where()
@@ -691,6 +712,14 @@ class Table(Queryable):
             else " ({})".format(", ".join(c.name for c in self.columns)),
         )
 
+    @property
+    def count(self):
+        if self.db.use_counts_table:
+            counts = self.db.cached_counts([self.name])
+            if counts:
+                return next(iter(counts.values()))
+        return self.execute_count()
+
     def exists(self):
         return self.name in self.db.table_names()
 
@@ -1227,7 +1256,9 @@ class Table(Queryable):
             )
             .strip()
             .format(
-                create_counts_table=self.db._counts_table_create,
+                create_counts_table=_COUNTS_TABLE_CREATE_SQL.format(
+                    self.db._counts_table_name
+                ),
                 counts_table=self.db._counts_table_name,
                 table=self.name,
                 table_quoted=self.db.quote(self.name),
@@ -1235,6 +1266,7 @@ class Table(Queryable):
         )
         with self.db.conn:
             self.db.conn.executescript(sql)
+        self.db.use_counts_table = True
 
     def enable_fts(
         self,
@@ -1650,6 +1682,7 @@ class Table(Queryable):
         )
 
         with self.db.conn:
+            result = None
             for query, params in queries_and_params:
                 try:
                     result = self.db.execute(query, params)

tests/test_enable_counts.py

@@ -96,3 +96,36 @@ def test_cli_enable_counts(counts_db_path, extra_args, expected_triggers):
     result = CliRunner().invoke(cli.cli, ["enable-counts", counts_db_path] + extra_args)
     assert result.exit_code == 0
     assert list(db.triggers_dict.keys()) == expected_triggers
+
+
+def test_uses_counts_after_enable_counts(counts_db_path):
+    db = Database(counts_db_path)
+    logged = []
+    with db.tracer(lambda sql, parameters: logged.append((sql, parameters))):
+        assert db["foo"].count == 1
+        assert logged == [
+            ("select name from sqlite_master where type = 'view'", None),
+            ("select count(*) from [foo]", None),
+        ]
+        logged.clear()
+        assert not db.use_counts_table
+        db.enable_counts()
+        assert db.use_counts_table
+        assert db["foo"].count == 1
+    assert logged == [
+        (
+            "CREATE TABLE IF NOT EXISTS [_counts](\n   [table] TEXT PRIMARY KEY,\n   count INTEGER DEFAULT 0\n);",
+            None,
+        ),
+        ("select name from sqlite_master where type = 'table'", None),
+        ("select name from sqlite_master where type = 'view'", None),
+        ("select name from sqlite_master where type = 'view'", None),
+        ("select name from sqlite_master where type = 'view'", None),
+        ("select sql from sqlite_master where name = ?", ("foo",)),
+        ("SELECT quote(:value)", {"value": "foo"}),
+        ("select sql from sqlite_master where name = ?", ("bar",)),
+        ("SELECT quote(:value)", {"value": "bar"}),
+        ("select sql from sqlite_master where name = ?", ("_counts",)),
+        ("select name from sqlite_master where type = 'view'", None),
+        ("select [table], count from _counts where [table] in (?)", ["foo"]),
+    ]