1"""
  2Database test helpers.
  3"""
  4
  5from collections.abc import Generator
  6from contextlib import contextmanager
  7from dataclasses import dataclass
  8from functools import partial
  9from typing import Any
 10from weakref import WeakKeyDictionary
 11
 12from plain.testing import Captured, CaptureSource
 13
 14from ..connection import DatabaseConnection
 15from ..db import get_connection
 16
 17__all__ = ["CapturedQueries", "CapturedQuery", "capture_queries", "max_queries"]
 18
 19
 20@dataclass(frozen=True)
 21class CapturedQuery:
 22    """One thing the database was asked to do during a `capture_queries` block."""
 23
 24    # The statement as it was sent: `%s` placeholders in place, values apart.
 25    sql: str
 26    # The same statement with the values filled in. An `executemany()` is one
 27    # query here, prefixed with how many times it ran: "3 times: INSERT ...".
 28    sql_with_params: str
 29    # False for BEGIN, COMMIT and ROLLBACK, which the connection issues
 30    # itself. A savepoint is a statement like any other.
 31    is_statement: bool
 32
 33
 34class CapturedQueries(Captured[CapturedQuery]):
 35    """
 36    What the database was asked to do during a `capture_queries` block, in
 37    the order it was asked.
 38    """
 39
 40    def __init__(self) -> None:
 41        super().__init__(helper="capture_queries")
 42
 43    def sql_statements(self, *, table: str | None = None) -> list[str]:
 44        """
 45        The statements as they were sent, with runs of whitespace collapsed
 46        to single spaces — the form to compare against SQL written out in a
 47        test.
 48
 49            with capture_queries() as queries:
 50                cache.set_many({"a": 1, "b": 2})
 51
 52            [statement] = queries.sql_statements()
 53            assert statement.startswith('INSERT INTO "plaincache_cacheditem"')
 54
 55        Savepoints are statements and are here. BEGIN, COMMIT and ROLLBACK
 56        are not: the connection issues those itself.
 57
 58        Pass `table` to keep only the statements that name that table as a
 59        quoted identifier: `table="users_user"` keeps the statements that
 60        contain `"users_user"`.
 61        """
 62        statements = []
 63        for query in self:
 64            if not query.is_statement:
 65                continue
 66            statement = " ".join(query.sql.split())
 67            if table is not None and f'"{table}"' not in statement:
 68                continue
 69            statements.append(statement)
 70        return statements
 71
 72
 73@contextmanager
 74def capture_queries() -> Generator[CapturedQueries]:
 75    """
 76    What the database is asked to do within the block.
 77
 78        with capture_queries() as queries:
 79            list(Article.query.all())
 80
 81        assert len(queries) == 1
 82        assert queries[0].sql.startswith("SELECT")
 83
 84    It records at the connection, so it sees every query whether or not the
 85    query is traced, and whether or not the connection was returned to the
 86    pool and taken out again along the way, as it is by a request outside a
 87    transaction.
 88    """
 89    conn = get_connection()
 90    previous = conn.force_debug_cursor
 91    conn.force_debug_cursor = True
 92    captured = CapturedQueries()
 93    try:
 94        with _query_source_of(conn).capturing_into(captured):
 95            yield captured
 96    finally:
 97        conn.force_debug_cursor = previous
 98
 99
100# Every capture on a connection reads that connection's query log, so each
101# connection has one source, kept for as long as the connection is.
102_query_sources: WeakKeyDictionary[DatabaseConnection, CaptureSource[CapturedQuery]] = (
103    WeakKeyDictionary()
104)
105
106
107def _query_source_of(conn: DatabaseConnection) -> CaptureSource[CapturedQuery]:
108    source = _query_sources.get(conn)
109    if source is None:
110        # The source holds the connection's log, not the connection, so it
111        # doesn't keep a closed connection from being let go.
112        source = CaptureSource(
113            read=partial(_queries_in, conn.captured_queries_log),
114            clear=conn.captured_queries_log.clear,
115        )
116        _query_sources[conn] = source
117    return source
118
119
120def _queries_in(queries_log: list[dict[str, Any]]) -> list[CapturedQuery]:
121    return [
122        CapturedQuery(
123            sql=entry.get("sql_as_sent", entry["sql"]),
124            sql_with_params=entry["sql"],
125            is_statement="sql_as_sent" in entry,
126        )
127        for entry in queries_log
128    ]
129
130
131@contextmanager
132def max_queries(count: int) -> Generator[None]:
133    """
134    Fail if more than `count` database queries execute within the block.
135
136        with max_queries(5):
137            client.get("/dashboard/")
138
139    A query budget is a contract, checked on every run.
140    """
141    with capture_queries() as queries:
142        yield
143    executed = len(queries)
144    if executed > count:
145        sql_lines = "\n".join(f"  {query.sql_with_params}" for query in queries)
146        raise AssertionError(
147            f"Expected at most {count} queries, {executed} were executed:\n{sql_lines}"
148        )