diff --git a/docs/cli.rst b/docs/cli.rst index 3a55629..46ad820 100644 --- a/docs/cli.rst +++ b/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. diff --git a/docs/python-api.rst b/docs/python-api.rst index 04703de..b02fafa 100644 --- a/docs/python-api.rst +++ b/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 ================ diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 719d03c..f1f3de6 100644 --- a/sqlite_utils/db.py +++ b/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) diff --git a/tests/test_enable_counts.py b/tests/test_enable_counts.py index a6f35a0..6fabd78 100644 --- a/tests/test_enable_counts.py +++ b/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"]), + ]