diff --git a/docs/agents.md b/docs/agents.md index 61ec994..0547032 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -156,7 +156,7 @@ from the submodule named beside it further down. | Base ORM | `Base`, `BaseTable`, `DatetimeColumnsMixin`, `DB_NAMING_CONVENTION`, `PydanticJSONB`, `GenericJSONDict`, `UnConstrainedEnum`, `load_orm_metadata` | | Engine utilities | `build_engine_kwargs`, `resolve_pool_class`, `register_pool_class`, `PoolRegistry`, `PoolClassStr`, `configure_orjson_serialization` | | Config protocols | `PostgresSettingsProtocol`, `ConnectionSettingsProtocol`, `PoolSettingsProtocol`, `QuerySettingsProtocol` | -| Metrics protocols | `PostgresMetricsProtocol`, `PoolStatsRecorder`, `CheckoutRecorder`, `CheckoutWaitRecorder`, `ErrorRecorder` | +| Metrics protocols | `PostgresMetricsProtocol`, `PoolStatsRecorder`, `PoolStatsFunctionRecorder`, `CheckoutRecorder`, `CheckoutWaitRecorder`, `ErrorRecorder` | | Version | `__version__` | ### `AsyncSessionManager` @@ -202,7 +202,13 @@ are startup parameters, so read rule 16 before adding one. `sqlalchemy_foundation_kit.session.manager`, not re-exported — wire the pool listeners, respectively the per-transaction `search_path`, onto an engine you built yourself. The manager calls them when `metrics` / `search_path` is passed. A raising metrics callback is -logged and swallowed, never propagated. +logged and swallowed, never propagated. The pool gauges are read when the registry is +collected if `metrics` implements `PoolStatsFunctionRecorder` +(`set_pool_stats_functions(pool_size, pool_checked_out, pool_overflow)`, each a +zero-argument callable over the engine's current pool), as `PostgresMetrics` does; +otherwise `record_pool_stats` is called from the `checkout` / `checkin` events and lags one +return behind, so an idle pool keeps its last burst. One `PostgresMetrics` reports one +engine. `instrument_pool_class(pool_class, metrics)` lives beside them and is the other half of the metrics wiring: it returns a subclass of `pool_class` that times `Pool.connect()`, which is @@ -336,9 +342,9 @@ variables only through the parent `BaseSettings` that holds it. Under | Series | Type | What it measures | |---|---|---| -| `postgres_db_pool_size` | gauge | connections the pool holds | +| `postgres_db_pool_size` | gauge | the configured `pool_size`; 0 for `null` / `static` | | `postgres_db_pool_checked_out` | gauge | connections currently in use | -| `postgres_db_pool_overflow` | gauge | connections over `pool_size`, within `max_overflow` | +| `postgres_db_pool_overflow` | gauge | open connections over `pool_size`, within `max_overflow`; never negative | | `postgres_db_connection_checkout_wait_seconds` | histogram | how long a caller **waited** for a connection — the queue wait, plus the pre-ping and the connect handshake when the pool had to grow. Buckets run to 30 s, the default `pool_timeout` | | `postgres_db_connection_held_duration_seconds` | histogram | how long a caller **held** one, checkout to checkin — the query time seen from the pool | | `postgres_db_connection_checkout_duration_seconds` | histogram | **deprecated**, and never measured what its name says: it is the held duration under its old name. Kept for one more minor release, then dropped. Move dashboards to `…_held_duration_seconds`, or to `…_checkout_wait_seconds` if what you wanted was the wait | @@ -446,7 +452,7 @@ decorator. ### Fixed since 0.2.0 Three published entry points raise before they reach the database on 0.2.0; on 0.2.1 the -PgBouncer-safe defaults cannot connect through PgBouncer; and up to 0.3.0 the pool metrics +PgBouncer-safe defaults cannot connect through PgBouncer; and up to 0.4.0 the pool metrics do not say what their names say. All are right on current versions; the workaround column is what to do on the version the row names. @@ -457,6 +463,7 @@ is what to do on the version the row names. | `import sqlalchemy_foundation_kit.contrib.di` (or `.contrib.dependency_injector`) without its extra (0.2.0) | `AttributeError: 'NoneType' object has no attribute 'APP'` (resp. `'DeclarativeContainer'`) instead of the intended `ImportError` | install the extra; the message is not the one the code meant to give you | | `create_async_session_manager(config)` through PgBouncer in transaction mode (0.2.1 and earlier) | `jit="off"` was the default and `db_schema` went as `search_path`, both as startup parameters: a default PgBouncer refuses every connection with `ProtocolViolationError: unsupported startup parameter: jit`; with `ignore_startup_parameters=jit,search_path` it connects and silently drops both, so every query lands in `public` | `jit=None`, `db_schema=None`, and `ALTER ROLE … SET search_path` on the server | | `postgres_db_connection_timeouts_total` (0.3.0 and earlier) | stayed at zero through every pool checkout timeout. The counter was fed only from the engine's `handle_error` listener, and a pool `TimeoutError` is raised by `pool.connect()` before any DBAPI call, so it never reaches that listener — the one timeout a pool actually produces under load was the one the counter did not count | count `sqlalchemy.exc.TimeoutError` around your own session calls | +| `postgres_db_pool_overflow`, `postgres_db_pool_checked_out` (0.4.0 and earlier) | sampled from the `checkout` / `checkin` events: `overflow` exported `QueuePool.overflow()`, which starts at `-pool_size` (`-9` with one connection open of 10), and since `checkin` fires before the connection is back, an idle pool kept reporting its last burst as checked out | read `pool.checkedout()` and `max(pool.overflow(), 0)` yourself | | `postgres_db_connection_checkout_duration_seconds` (0.3.0 and earlier) | the only checkout histogram there was, and it measures the time between the `checkout` and `checkin` events — how long a connection was *held*, which is query duration seen from the pool. The time a caller waited for a connection, which is what the name suggests and what predicts a pool outage, was not exposed at all | none; the wait was not measurable from outside the pool | `IsolationLevel` itself was always fine — `READ_UNCOMMITTED`, `READ_COMMITTED`, diff --git a/docs/guide/advanced.md b/docs/guide/advanced.md index 3f1e659..b653b1d 100644 --- a/docs/guide/advanced.md +++ b/docs/guide/advanced.md @@ -268,9 +268,9 @@ session_manager = create_async_session_manager( | Metric | Type | Description | |--------|------|-------------| -| `myapp_postgres_db_pool_size` | Gauge | Current pool size | +| `myapp_postgres_db_pool_size` | Gauge | Configured `pool_size` (0 for a pool without a queue) | | `myapp_postgres_db_pool_checked_out` | Gauge | Connections currently in use | -| `myapp_postgres_db_pool_overflow` | Gauge | Overflow connections created | +| `myapp_postgres_db_pool_overflow` | Gauge | Open connections over `pool_size`, within `max_overflow` | | `myapp_postgres_db_connection_checkout_wait_seconds` | Histogram | Time a caller **waited** for a connection | | `myapp_postgres_db_connection_held_duration_seconds` | Histogram | Time a caller **held** a connection, checkout to checkin | | `myapp_postgres_db_connection_checkout_duration_seconds` | Histogram | Deprecated alias of the held duration — see below | @@ -279,6 +279,12 @@ session_manager = create_async_session_manager( Without a `prefix` the names are `postgres_db_pool_size` and so on. +The three gauges are read from the pool when Prometheus scrapes, not sampled from pool +events, so an idle pool reports `checked_out` 0 instead of whatever its last burst left +behind. A custom metrics object gets the same by implementing `PoolStatsFunctionRecorder`; +one that only implements `record_pool_stats` is sampled on checkout and checkin, and lags +one return behind. + **Wait or held — they are not the same number.** The wait is the time spent inside `pool.connect()` before a connection is handed out: the queue wait, plus the pre-ping and the connect handshake when the pool has to grow. It is the metric that predicts a pool @@ -333,7 +339,8 @@ async def metrics_endpoint(): **Grafana Dashboard Example:** ```promql -# Pool utilization +# Pool utilization against pool_size -- above 100 while overflow connections are in use; +# the hard ceiling is pool_size + max_overflow (myapp_postgres_db_pool_checked_out / myapp_postgres_db_pool_size) * 100 # P95 wait for a connection -- the one to alert on diff --git a/sqlalchemy_foundation_kit/__init__.py b/sqlalchemy_foundation_kit/__init__.py index 5dcde47..fe606ce 100644 --- a/sqlalchemy_foundation_kit/__init__.py +++ b/sqlalchemy_foundation_kit/__init__.py @@ -33,6 +33,7 @@ CheckoutRecorder, CheckoutWaitRecorder, ErrorRecorder, + PoolStatsFunctionRecorder, PoolStatsRecorder, PostgresMetricsProtocol, ) @@ -92,6 +93,7 @@ "CheckoutRecorder", "CheckoutWaitRecorder", "ErrorRecorder", + "PoolStatsFunctionRecorder", "PoolStatsRecorder", "PostgresMetricsProtocol", # Session Management diff --git a/sqlalchemy_foundation_kit/contrib/metrics/postgres.py b/sqlalchemy_foundation_kit/contrib/metrics/postgres.py index 95a9b8d..70c2f5b 100644 --- a/sqlalchemy_foundation_kit/contrib/metrics/postgres.py +++ b/sqlalchemy_foundation_kit/contrib/metrics/postgres.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from collections.abc import Callable try: from prometheus_client import Counter, Gauge, Histogram @@ -59,9 +60,9 @@ class PostgresMetrics: """Postgres connection pool metrics. Metrics: - - postgres_db_pool_size: Current database connection pool size. + - postgres_db_pool_size: Configured pool_size of the connection pool. - postgres_db_pool_checked_out: Number of connections currently checked out. - - postgres_db_pool_overflow: Number of connections over pool_size (within max_overflow). + - postgres_db_pool_overflow: Number of open connections over pool_size (within max_overflow). - postgres_db_connection_checkout_wait_seconds: Time a caller waited for a connection. - postgres_db_connection_held_duration_seconds: Time a connection was held by its caller. - postgres_db_connection_checkout_duration_seconds: Deprecated alias of the held duration, @@ -72,6 +73,10 @@ class PostgresMetrics: Labels: - error_type: Type of connection error (for errors_total). + The three pool gauges are read from the pool when the registry is collected, not + sampled from pool events, so an idle pool reports what it holds. One instance reports + one engine: attached to a second engine, the gauges follow the second. + The wait histogram and the timeout counter are fed by the instrumented pool class the session manager builds, so they only move on an engine created by ``AsyncSessionManager`` (or one whose pool went through ``instrument_pool_class``). @@ -90,7 +95,7 @@ def __init__(self, prefix: str | None = None) -> None: self.pool_size = Gauge( _make_metric_name("postgres_db_pool_size", prefix), - "Current database connection pool size", + "Configured pool_size of the database connection pool", ) self.pool_checked_out = Gauge( _make_metric_name("postgres_db_pool_checked_out", prefix), @@ -98,7 +103,7 @@ def __init__(self, prefix: str | None = None) -> None: ) self.pool_overflow = Gauge( _make_metric_name("postgres_db_pool_overflow", prefix), - "Number of connections over pool_size (within max_overflow)", + "Number of open connections over pool_size (within max_overflow)", ) self.connection_checkout_wait = Histogram( _make_metric_name("postgres_db_connection_checkout_wait_seconds", prefix), @@ -136,6 +141,17 @@ def record_pool_stats( self.pool_checked_out.set(float(pool_checked_out)) self.pool_overflow.set(float(pool_overflow)) + def set_pool_stats_functions( + self, + pool_size: Callable[[], float], + pool_checked_out: Callable[[], float], + pool_overflow: Callable[[], float], + ) -> None: + """Read the pool gauges from these functions whenever the registry is collected.""" + self.pool_size.set_function(pool_size) + self.pool_checked_out.set_function(pool_checked_out) + self.pool_overflow.set_function(pool_overflow) + def record_checkout( self, duration: float, diff --git a/sqlalchemy_foundation_kit/protocols/__init__.py b/sqlalchemy_foundation_kit/protocols/__init__.py index 4f786cf..84db69c 100644 --- a/sqlalchemy_foundation_kit/protocols/__init__.py +++ b/sqlalchemy_foundation_kit/protocols/__init__.py @@ -10,6 +10,7 @@ CheckoutRecorder, CheckoutWaitRecorder, ErrorRecorder, + PoolStatsFunctionRecorder, PoolStatsRecorder, PostgresMetricsProtocol, ) @@ -18,6 +19,7 @@ "CheckoutRecorder", "CheckoutWaitRecorder", "ErrorRecorder", + "PoolStatsFunctionRecorder", "PoolStatsRecorder", "PostgresMetricsProtocol", ] diff --git a/sqlalchemy_foundation_kit/protocols/metrics.py b/sqlalchemy_foundation_kit/protocols/metrics.py index 33467f9..166e4e7 100644 --- a/sqlalchemy_foundation_kit/protocols/metrics.py +++ b/sqlalchemy_foundation_kit/protocols/metrics.py @@ -7,6 +7,7 @@ from __future__ import annotations +from collections.abc import Callable from typing import Protocol, runtime_checkable @@ -21,10 +22,51 @@ def record_pool_stats( ) -> None: """Record database connection pool statistics. + Called from the pool's ``checkout`` and ``checkin`` events, so each call describes + the pool as that event saw it -- and ``checkin`` fires before the connection is + back, so the last call of a burst still counts it as checked out. A recorder that + also implements :class:`PoolStatsFunctionRecorder` reads the pool when it is + collected instead, and this method is then not called. + Args: - pool_size: Current total number of connections in the pool. + pool_size: The configured ``pool_size`` of a queue pool, 0 for any other pool. pool_checked_out: Number of connections currently in use. - pool_overflow: Number of connections over the configured pool_size. + pool_overflow: Number of open connections over ``pool_size``, within + ``max_overflow``; never negative. + """ + ... + + +@runtime_checkable +class PoolStatsFunctionRecorder(Protocol): + """Capability protocol for reading pool statistics when the metrics are collected. + + A gauge sampled from pool events goes stale: the ``checkin`` event fires before the + connection is returned, and nothing fires once the pool is idle, so an idle pool keeps + reporting the last burst. A recorder that implements this protocol is handed + functions that read the pool on demand -- ``prometheus_client``'s + ``Gauge.set_function`` -- and :class:`PoolStatsRecorder` is then not used for it. + + Not part of :class:`PostgresMetricsProtocol`, so implementations written without it + stay valid; the session manager tests for it with ``isinstance``. + """ + + def set_pool_stats_functions( + self, + pool_size: Callable[[], float], + pool_checked_out: Callable[[], float], + pool_overflow: Callable[[], float], + ) -> None: + """Read the pool statistics from these functions whenever they are collected. + + Each function reads the engine's current pool, so the pool that ``dispose()`` + creates is the one reported. Values mean the same as in + :meth:`PoolStatsRecorder.record_pool_stats`. + + Args: + pool_size: Returns the configured ``pool_size``. + pool_checked_out: Returns the number of connections currently in use. + pool_overflow: Returns the number of open connections over ``pool_size``. """ ... @@ -104,6 +146,7 @@ class PostgresMetricsProtocol(PoolStatsRecorder, CheckoutRecorder, ErrorRecorder "CheckoutRecorder", "CheckoutWaitRecorder", "ErrorRecorder", + "PoolStatsFunctionRecorder", "PoolStatsRecorder", "PostgresMetricsProtocol", ] diff --git a/sqlalchemy_foundation_kit/session/manager.py b/sqlalchemy_foundation_kit/session/manager.py index 4c6595e..5d1b2c0 100644 --- a/sqlalchemy_foundation_kit/session/manager.py +++ b/sqlalchemy_foundation_kit/session/manager.py @@ -21,10 +21,11 @@ from .._typing import SessionT from ..base import build_engine_kwargs, resolve_pool_class -from ..protocols import CheckoutWaitRecorder +from ..protocols import CheckoutWaitRecorder, PoolStatsFunctionRecorder if TYPE_CHECKING: from sqlalchemy.engine import Connection + from sqlalchemy.pool import Pool from ..config import PoolSettingsProtocol from ..protocols import PostgresMetricsProtocol @@ -54,6 +55,25 @@ def _safe_metric_call(func: Callable[[], None], error_msg: str) -> None: logger.exception(error_msg) +def _pool_size(pool: Pool) -> int: + """The configured ``pool_size`` of a queue pool, 0 for a pool without one.""" + return int(pool.size()) if hasattr(pool, "size") else 0 + + +def _pool_checked_out(pool: Pool) -> int: + """Connections currently in use.""" + return int(pool.checkedout()) if hasattr(pool, "checkedout") else 0 + + +def _pool_overflow(pool: Pool) -> int: + """Open connections over ``pool_size``. + + ``QueuePool.overflow()`` starts at ``-pool_size`` and counts every connection the pool + opens, so it is negative until the pool has grown past ``pool_size``. + """ + return max(int(pool.overflow()), 0) if hasattr(pool, "overflow") else 0 + + def instrument_pool_class(pool_class: type, metrics: CheckoutWaitRecorder) -> type: """Build a pool class that times the wait for a connection and counts pool timeouts. @@ -121,6 +141,12 @@ def attach_metrics(engine: AsyncEngine, metrics: PostgresMetricsProtocol) -> Non Registers event handlers for connection checkout, checkin, and error events to collect pool statistics and connection metrics. + The pool gauges come from the pool itself: a ``metrics`` that implements + ``PoolStatsFunctionRecorder`` is handed functions that read ``engine.pool`` whenever + the metrics are collected, so an idle pool reports what it holds rather than its last + burst, and a pool recreated by ``dispose()`` is the one read. One that only implements + ``record_pool_stats`` is sampled from the ``checkout`` and ``checkin`` events instead. + The checkout wait and the pool checkout timeout are **not** listeners — no pool event fires for either — and come from the pool class :func:`instrument_pool_class` builds. If ``metrics`` records the wait but the engine's pool was built without it, that is @@ -142,12 +168,24 @@ def attach_metrics(engine: AsyncEngine, metrics: PostgresMetricsProtocol) -> Non type(pool).__name__, ) + reads_pool_on_collect = False + if isinstance(metrics, PoolStatsFunctionRecorder): + reads_pool_on_collect = True + # engine.pool is looked up on every read: dispose() replaces the pool object. + metrics.set_pool_stats_functions( + pool_size=lambda: _pool_size(engine.pool), + pool_checked_out=lambda: _pool_checked_out(engine.pool), + pool_overflow=lambda: _pool_overflow(engine.pool), + ) + def record_pool_stats() -> None: + if reads_pool_on_collect: + return _safe_metric_call( lambda: metrics.record_pool_stats( - pool_size=pool.size() if hasattr(pool, "size") else 0, - pool_checked_out=(pool.checkedout() if hasattr(pool, "checkedout") else 0), - pool_overflow=pool.overflow() if hasattr(pool, "overflow") else 0, + pool_size=_pool_size(pool), + pool_checked_out=_pool_checked_out(pool), + pool_overflow=_pool_overflow(pool), ), "Failed to record database pool stats", ) diff --git a/tests/integration/test_metrics_integration.py b/tests/integration/test_metrics_integration.py index daa7859..72f390c 100644 --- a/tests/integration/test_metrics_integration.py +++ b/tests/integration/test_metrics_integration.py @@ -203,3 +203,80 @@ def _exposed(prefix: str) -> dict[str, float]: name, value = line.rsplit(" ", 1) samples[name] = float(value) return samples + + +# ============================================================================ +# Pool Gauges Tests +# ============================================================================ + + +@dataclass +class _GrowingPoolSettings: + """Ten connections that can grow by twenty: the defaults, with a pool to watch.""" + + kind: str = "async_adapted_queue" + size: int = 10 + max_overflow: int = 20 + pre_ping: bool = False + recycle: int = -1 + timeout: float = 5.0 + + +async def _while_holding(manager: AsyncSessionManager[AsyncSession], n: int, prefix: str) -> dict[str, float]: + """Check out ``n`` connections at once, read the registry while all are held, release them.""" + all_checked_out, release = asyncio.Event(), asyncio.Event() + ready = 0 + + async def one() -> None: + nonlocal ready + async with manager.get_session() as session: + await session.execute(text("SELECT 1")) + ready += 1 + if ready == n: + all_checked_out.set() + await release.wait() + + tasks = [asyncio.create_task(one()) for _ in range(n)] + await all_checked_out.wait() + exposed = _exposed(prefix) + release.set() + await asyncio.gather(*tasks) + return exposed + + +def _gauges(exposed: dict[str, float], prefix: str) -> tuple[float, float, float]: + """``(pool_size, pool_checked_out, pool_overflow)`` as scraped.""" + return ( + exposed[f"{prefix}_postgres_db_pool_size"], + exposed[f"{prefix}_postgres_db_pool_checked_out"], + exposed[f"{prefix}_postgres_db_pool_overflow"], + ) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__postgres_metrics__pool_gauges__describe_the_pool_while_busy_and_once_idle( + postgres_container: PostgresContainer, +) -> None: + # Arrange: the reporter's run -- size=10, max_overflow=20, one request, then a burst of 12 + pytest.importorskip("prometheus_client") + from sqlalchemy_foundation_kit.contrib.metrics import PostgresMetrics + + prefix = f"it_{uuid.uuid4().hex[:8]}" + manager: AsyncSessionManager[AsyncSession] = AsyncSessionManager( + asyncpg_url(postgres_container), + poolclass="async_adapted_queue", + metrics=PostgresMetrics(prefix=prefix), + pool_settings=_GrowingPoolSettings(), + ) + + async with manager: + # Act & Assert: one connection in use is not "-9 overflow" + assert _gauges(await _while_holding(manager, 1, prefix), prefix) == (10.0, 1.0, 0.0) + # ...and once it is back, nothing is in use + assert _gauges(_exposed(prefix), prefix) == (10.0, 0.0, 0.0) + + # Act & Assert: twelve at once is two past pool_size + assert _gauges(await _while_holding(manager, 12, prefix), prefix) == (10.0, 12.0, 2.0) + # ...and after the burst the overflow connections are closed and nothing is in use + assert _gauges(_exposed(prefix), prefix) == (10.0, 0.0, 0.0) diff --git a/tests/unit/contrib/metrics/test_postgres.py b/tests/unit/contrib/metrics/test_postgres.py index 8c6ec24..231ebb0 100644 --- a/tests/unit/contrib/metrics/test_postgres.py +++ b/tests/unit/contrib/metrics/test_postgres.py @@ -31,6 +31,11 @@ def _histogram(histogram: object, suffix: str) -> float: raise AssertionError(f"no {suffix} sample on the histogram") +def _collected(gauge: object) -> float: + """The value a scrape of this gauge would expose right now.""" + return float(next(iter(gauge.collect())).samples[0].value) # type: ignore[attr-defined] + + # ============================================================================ # _make_metric_name Tests # ============================================================================ @@ -187,6 +192,36 @@ def test__postgres_metrics__record_pool_stats__sets_gauge_values() -> None: assert metrics.pool_overflow._value.get() == 2.0 # type: ignore[attr-defined] +@pytest.mark.unit +def test__postgres_metrics__set_pool_stats_functions__gauges_read_them_on_collect() -> None: + # Arrange + metrics = PostgresMetrics(prefix=_unique_prefix()) + checked_out = [3] + + # Act + metrics.set_pool_stats_functions( + pool_size=lambda: 10, + pool_checked_out=lambda: checked_out[0], + pool_overflow=lambda: 0, + ) + busy = _collected(metrics.pool_checked_out) + checked_out[0] = 0 + idle = _collected(metrics.pool_checked_out) + + # Assert + assert _collected(metrics.pool_size) == 10.0 + assert (busy, idle) == (3.0, 0.0) + + +@pytest.mark.unit +def test__postgres_metrics__satisfies_the_pool_stats_function_protocol() -> None: + # Arrange + from sqlalchemy_foundation_kit.protocols import PoolStatsFunctionRecorder + + # Act & Assert: attach_metrics picks the collect-time path by this check + assert isinstance(PostgresMetrics(prefix=_unique_prefix()), PoolStatsFunctionRecorder) + + @pytest.mark.unit def test__postgres_metrics__record_pool_stats__updates_existing_values() -> None: # Arrange diff --git a/tests/unit/session/test_manager.py b/tests/unit/session/test_manager.py index bffed54..3c19f59 100644 --- a/tests/unit/session/test_manager.py +++ b/tests/unit/session/test_manager.py @@ -4,6 +4,7 @@ import threading import time +from collections.abc import Callable from typing import cast from unittest.mock import AsyncMock, Mock, patch @@ -63,6 +64,15 @@ def test__safe_metric_call__exception__does_not_propagate() -> None: # ============================================================================ +def _event_sampled_metrics() -> Mock: + """Metrics without ``set_pool_stats_functions``, so the pool events sample the gauges. + + A bare ``Mock()`` will not do: before Python 3.12 a runtime protocol check is a + ``hasattr`` per member, which a ``Mock`` passes for any name. + """ + return Mock(spec=["record_pool_stats", "record_checkout", "record_error"]) + + def test__attach_metrics__registers_checkout_listener() -> None: # Arrange mock_engine = Mock() @@ -118,7 +128,7 @@ def test__attach_metrics__checkout_records_pool_stats() -> None: mock_engine = Mock() mock_engine.pool = mock_pool - mock_metrics = Mock() + mock_metrics = _event_sampled_metrics() # Act with patch("sqlalchemy_foundation_kit.session.manager.event") as mock_event: @@ -144,6 +154,125 @@ def test__attach_metrics__checkout_records_pool_stats() -> None: assert call_args["pool_overflow"] == 2 +def test__attach_metrics__checkout_on_a_pool_below_its_size__records_no_negative_overflow() -> None: + # Arrange: QueuePool.overflow() is "open connections minus pool_size" -- -9 with one open + mock_engine = Mock() + mock_engine.pool.size.return_value = 10 + mock_engine.pool.checkedout.return_value = 1 + mock_engine.pool.overflow.return_value = -9 + mock_metrics = _event_sampled_metrics() + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event") as mock_event: + attach_metrics(mock_engine, mock_metrics) + checkout_callback = next(c[0][2] for c in mock_event.listen.call_args_list if c[0][1] == "checkout") + checkout_callback(None, Mock(info={}), None) + + # Assert + assert mock_metrics.record_pool_stats.call_args[1]["pool_overflow"] == 0 + + +class _FunctionRecorder: + """Reads the pool gauges on demand, the way ``Gauge.set_function`` does.""" + + def __init__(self) -> None: + self.functions: dict[str, Callable[[], float]] = {} + self.event_samples = 0 + self.held: list[float] = [] + + def set_pool_stats_functions( + self, + pool_size: Callable[[], float], + pool_checked_out: Callable[[], float], + pool_overflow: Callable[[], float], + ) -> None: + self.functions = {"size": pool_size, "checked_out": pool_checked_out, "overflow": pool_overflow} + + def record_pool_stats(self, pool_size: int, pool_checked_out: int, pool_overflow: int) -> None: + self.event_samples += 1 + + def record_checkout(self, duration: float) -> None: + self.held.append(duration) + + def record_error(self, error_type: str, is_timeout: bool = False) -> None: ... + + def read(self) -> dict[str, float]: + return {name: function() for name, function in self.functions.items()} + + +def _queue_pool(size: int, checked_out: int, overflow: int) -> Mock: + pool = Mock() + pool.size.return_value = size + pool.checkedout.return_value = checked_out + pool.overflow.return_value = overflow + return pool + + +def test__attach_metrics__function_recorder__reads_the_pool_when_collected() -> None: + # Arrange + mock_engine = Mock() + mock_engine.pool = _queue_pool(size=10, checked_out=1, overflow=-9) + recorder = _FunctionRecorder() + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event"): + attach_metrics(mock_engine, recorder) # type: ignore[arg-type] + busy = recorder.read() + mock_engine.pool.checkedout.return_value = 0 # the connection came back; no event follows + idle = recorder.read() + + # Assert + assert busy == {"size": 10, "checked_out": 1, "overflow": 0} + assert idle == {"size": 10, "checked_out": 0, "overflow": 0} + + +def test__attach_metrics__function_recorder__follows_the_pool_dispose_creates() -> None: + # Arrange + mock_engine = Mock() + mock_engine.pool = _queue_pool(size=10, checked_out=3, overflow=-7) + recorder = _FunctionRecorder() + with patch("sqlalchemy_foundation_kit.session.manager.event"): + attach_metrics(mock_engine, recorder) # type: ignore[arg-type] + + # Act: engine.dispose() swaps in a fresh pool + mock_engine.pool = _queue_pool(size=10, checked_out=0, overflow=-10) + + # Assert + assert recorder.read() == {"size": 10, "checked_out": 0, "overflow": 0} + + +def test__attach_metrics__function_recorder__pool_events_do_not_sample_the_gauges() -> None: + # Arrange + mock_engine = Mock() + recorder = _FunctionRecorder() + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event") as mock_event: + attach_metrics(mock_engine, recorder) # type: ignore[arg-type] + callbacks = {c[0][1]: c[0][2] for c in mock_event.listen.call_args_list} + record = Mock(info={}) + callbacks["checkout"](None, record, None) + callbacks["checkin"](None, record) + + # Assert: the held duration still comes from the events + assert recorder.event_samples == 0 + assert len(recorder.held) == 1 + + +def test__attach_metrics__function_recorder__pool_without_queue__reads_zero() -> None: + # Arrange + mock_engine = Mock() + mock_engine.pool = Mock(spec=[]) # NullPool / StaticPool have no size, checkedout, overflow + recorder = _FunctionRecorder() + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event"): + attach_metrics(mock_engine, recorder) # type: ignore[arg-type] + + # Assert + assert recorder.read() == {"size": 0, "checked_out": 0, "overflow": 0} + + def test__attach_metrics__checkout_stores_start_time() -> None: # Arrange mock_engine = Mock() @@ -281,7 +410,7 @@ def test__attach_metrics__pool_without_size__uses_zero() -> None: mock_pool = Mock(spec=[]) # No size method mock_engine = Mock() mock_engine.pool = mock_pool - mock_metrics = Mock() + mock_metrics = _event_sampled_metrics() # Act with patch("sqlalchemy_foundation_kit.session.manager.event") as mock_event: