datasette/tests/connection_budget.py

"""Record which databases a block of code touches, and assert a budget.

With a capped, idle-reaped connection pool, anything that loops over every
attached database opens (and evicts) a connection per database at a 0% hit
rate. These helpers make "this request only touches the databases it is
about" a testable property::

    async with connection_budget() as budget:
        await ds.client.get("/db1/t1.json")
    budget.assert_only("db1")  # plus _internal, always allowed

What counts as touching a database:

* ``opened``: ``Database.connect()`` calls - pooled read connections, write
  connections, isolated connections and the SchemaWatcher's short-lived
  scan/pragma connections (``track=False``)
* ``reads``: ``Database._execute_fn()`` calls (every ``execute()``,
  ``execute_fn()`` and introspection helper), whether or not the pool had an
  idle connection - so a warm pool does not hide a loop over every database
* ``writes``: ``execute_write*()``, ``execute_isolated_fn()`` and work run on
  the write connection
* ``attached``: databases ATTACHed to a connection while it was prepared
  (``--crossdb``)
* ``probes``: ``sqlite3.connect()`` calls that bypass ``Database.connect()``
  (``:memory:`` feature probes and the like), keyed by filename

The patches are installed process-wide for the duration of the block, so
background work (the SchemaWatcher's polling task, the read pool's reaper)
that runs during the block is counted too. Tests that want to measure one
request should turn polling off (``schema_watch_interval_ms=0``) or measure
it separately.
"""

import collections
import contextlib
import threading

from datasette import app as app_module
from datasette.database import Database
from datasette.utils import sqlite3

INTERNAL = "__INTERNAL__"

_local = threading.local()


class ConnectionBudget:
    def __init__(self):
        self._lock = threading.Lock()
        self.opened = collections.Counter()
        self.reads = collections.Counter()
        self.writes = collections.Counter()
        self.attached = collections.Counter()
        self.probes = collections.Counter()

    def _add(self, counter, key):
        with self._lock:
            counter[key] += 1

    def touched(self):
        """Names of every database opened, read, written or ATTACHed."""
        names = set()
        for counter in (self.opened, self.reads, self.writes, self.attached):
            names.update(counter)
        return names

    def summary(self):
        return {
            "opened": dict(self.opened),
            "reads": dict(self.reads),
            "writes": dict(self.writes),
            "attached": dict(self.attached),
            "probes": dict(self.probes),
        }

    def assert_only(self, *names, internal=True, probes=0):
        """Assert that nothing outside ``names`` was touched.

        ``_internal`` is allowed unless ``internal=False``. ``probes`` is the
        number of direct ``sqlite3.connect()`` calls allowed.
        """
        allowed = set(names)
        if internal:
            allowed.add(INTERNAL)
        extra = self.touched() - allowed
        assert not extra, (
            f"touched {len(extra)} database(s) outside the budget "
            f"{sorted(allowed)}: {sorted(extra)[:10]}"
            f"{' ...' if len(extra) > 10 else ''}\n{self.summary()}"
        )
        total_probes = sum(self.probes.values())
        assert (
            total_probes <= probes
        ), f"{total_probes} direct sqlite3.connect() call(s): {dict(self.probes)}"

    def assert_at_most(self, n, *, internal=True):
        """Assert that at most ``n`` user databases were touched."""
        touched = self.touched()
        if internal:
            touched.discard(INTERNAL)
        assert len(touched) <= n, (
            f"touched {len(touched)} databases, budget {n}: "
            f"{sorted(touched)[:10]}\n{self.summary()}"
        )

    def assert_no_connections(self, internal=True):
        """Assert that no connection was opened at all (``_internal`` aside)."""
        opened = {
            k: v for k, v in self.opened.items() if not (internal and k == INTERNAL)
        }
        assert not opened, f"opened connections: {opened}"
        assert not self.attached, f"attached: {dict(self.attached)}"
        assert not self.probes, f"probes: {dict(self.probes)}"


@contextlib.contextmanager
def _patched(budget):
    original_connect = Database.connect
    original_execute_fn = Database._execute_fn
    original_execute_write_fn = Database._execute_write_fn
    original_isolated = Database.execute_isolated_fn
    original_on_write_conn = Database._execute_on_write_connection
    original_prepare = app_module.Datasette._prepare_connection
    original_sqlite_connect = sqlite3.connect

    def connect(self, *args, **kwargs):
        budget._add(budget.opened, self.name)
        _local.depth = getattr(_local, "depth", 0) + 1
        try:
            return original_connect(self, *args, **kwargs)
        finally:
            _local.depth -= 1

    async def _execute_fn(self, fn):
        budget._add(budget.reads, self.name)
        return await original_execute_fn(self, fn)

    async def _execute_write_fn(self, fn, *args, **kwargs):
        budget._add(budget.writes, self.name)
        return await original_execute_write_fn(self, fn, *args, **kwargs)

    async def execute_isolated_fn(self, fn):
        budget._add(budget.writes, self.name)
        return await original_isolated(self, fn)

    async def _execute_on_write_connection(self, fn):
        budget._add(budget.writes, self.name)
        return await original_on_write_conn(self, fn)

    def _prepare_connection(self, conn, database, **kwargs):
        result = original_prepare(self, conn, database, **kwargs)
        try:
            rows = conn.execute("PRAGMA database_list").fetchall()
        except Exception:  # noqa: BLE001
            rows = []
        for row in rows:
            name = row[1]
            if name not in ("main", "temp"):
                budget._add(budget.attached, name)
        return result

    def sqlite_connect(database, *args, **kwargs):
        if not getattr(_local, "depth", 0):
            budget._add(budget.probes, str(database))
        return original_sqlite_connect(database, *args, **kwargs)

    Database.connect = connect
    Database._execute_fn = _execute_fn
    Database._execute_write_fn = _execute_write_fn
    Database.execute_isolated_fn = execute_isolated_fn
    Database._execute_on_write_connection = _execute_on_write_connection
    app_module.Datasette._prepare_connection = _prepare_connection
    sqlite3.connect = sqlite_connect
    try:
        yield budget
    finally:
        Database.connect = original_connect
        Database._execute_fn = original_execute_fn
        Database._execute_write_fn = original_execute_write_fn
        Database.execute_isolated_fn = original_isolated
        Database._execute_on_write_connection = original_on_write_conn
        app_module.Datasette._prepare_connection = original_prepare
        sqlite3.connect = original_sqlite_connect


class connection_budget:
    """Context manager (sync or async) yielding a ConnectionBudget."""

    def __init__(self):
        self.budget = ConnectionBudget()
        self._cm = None

    def __enter__(self):
        self._cm = _patched(self.budget)
        return self._cm.__enter__()

    def __exit__(self, *exc):
        return self._cm.__exit__(*exc)

    async def __aenter__(self):
        return self.__enter__()

    async def __aexit__(self, *exc):
        return self.__exit__(*exc)