diff --git a/README.md b/README.md index f0c5535..03021a7 100644 --- a/README.md +++ b/README.md @@ -243,7 +243,7 @@ from sqlalchemy_foundation_kit.contrib.settings import ( from sqlalchemy_foundation_kit.contrib.metrics import PostgresMetrics ``` - `PostgresMetrics` — Prometheus metrics for connection pool -- Tracks: pool size, checked out connections, checkout duration, errors +- Tracks: pool size, checked out connections, checkout wait, held duration, timeouts, errors #### `contrib.di` (requires `[dishka]`) ```python diff --git a/docs/agents.md b/docs/agents.md index 9bb3467..61ec994 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`, `ErrorRecorder` | +| Metrics protocols | `PostgresMetricsProtocol`, `PoolStatsRecorder`, `CheckoutRecorder`, `CheckoutWaitRecorder`, `ErrorRecorder` | | Version | `__version__` | ### `AsyncSessionManager` @@ -204,6 +204,15 @@ respectively the per-transaction `search_path`, onto an engine you built yoursel manager calls them when `metrics` / `search_path` is passed. A raising metrics callback is logged and swallowed, never propagated. +`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 +the only place the wait for a connection and the pool checkout timeout are visible — no +pool event fires for either. The manager applies it to whatever `resolve_pool_class` +returned whenever `metrics` implements `CheckoutWaitRecorder`, so callers get it without +asking. On an engine you build yourself, pass the result as `poolclass=` to +`create_async_engine`; `attach_metrics` logs a warning if you did not, rather than leaving +a wait histogram that never moves. + ### The unit of work `AsyncSQLAlchemyUnitOfWork(session_maker, transaction_factory, *, flush_before_commit=True)`. @@ -311,7 +320,7 @@ variables only through the parent `BaseSettings` that holds it. Under | Import | Needs | What you get | |---|---|---| -| `contrib.metrics.PostgresMetrics(prefix=None)` | `metrics` | the six gauges/histogram/counters below; pass it as `metrics=` to a manager | +| `contrib.metrics.PostgresMetrics(prefix=None)` | `metrics` | the eight series below; pass it as `metrics=` to a manager | | `contrib.telemetry.instrument_engine(engine, **kw)` | `telemetry` | the `on_engine_created` hook shape — traces one engine | | `contrib.telemetry.instrument_sqlalchemy(engine=None, **kw)` | `telemetry` | `SQLAlchemyInstrumentor().instrument(...)`; `engine=None` means every engine | | `contrib.telemetry.instrument_asyncpg(**kw)` | `telemetry` | `AsyncPGInstrumentor().instrument(...)` | @@ -323,11 +332,24 @@ variables only through the parent `BaseSettings` that holds it. Under | `contrib.dependency_injector.AsyncDatabaseResourceProvider(config, metrics=None, ...)` | `dependency-injector` | `await .start()` / `await .stop()` for a lifecycle you drive yourself | | `contrib.dependency_injector.PrometheusMetricsContainer` | `dependency-injector` + `metrics` | `postgres_metrics` from `metrics_settings`, `default_prefix`, `postgres_settings` | -`PostgresMetrics` publishes `postgres_db_pool_size`, `postgres_db_pool_checked_out`, -`postgres_db_pool_overflow` (gauges), `postgres_db_connection_checkout_duration_seconds` -(histogram), `postgres_db_connection_timeouts_total` and -`postgres_db_connection_errors_total{error_type}` (counters). A `prefix` is prepended with -an underscore and must match `^[a-zA-Z_][a-zA-Z0-9_]*$`. +`PostgresMetrics` publishes: + +| Series | Type | What it measures | +|---|---|---| +| `postgres_db_pool_size` | gauge | connections the pool holds | +| `postgres_db_pool_checked_out` | gauge | connections currently in use | +| `postgres_db_pool_overflow` | gauge | connections over `pool_size`, within `max_overflow` | +| `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 | +| `postgres_db_connection_timeouts_total` | counter | pool checkout timeouts, plus `TimeoutError`s seen by `handle_error` | +| `postgres_db_connection_errors_total{error_type}` | counter | database errors by exception class name | + +A `prefix` is prepended with an underscore and must match `^[a-zA-Z_][a-zA-Z0-9_]*$`. + +The wait histogram and the pool-timeout half of the counter come from the instrumented pool +class, not from a listener, so they move only on an engine built by `AsyncSessionManager` +or one whose pool went through `instrument_pool_class`. `retry_async_connection(connect_func, service_name, config=DEFAULT_RETRY_CONFIG)` is the startup retry the DI providers use, and is usable on its own: it awaits `connect_func()` @@ -412,12 +434,21 @@ decorator. connection leaks to the next client through a pooler; never issue one in a `connect` listener. Statement caches stay at 0 and `AsyncCConnection` stays, for PgBouncer before 1.22 (`max_prepared_statements=0`); 1.22+ tracks prepared statements itself. +17. **The wait for a connection and the time it was held are different metrics.** Alert on + `postgres_db_connection_checkout_wait_seconds` — a rising wait is a pool about to run + out, and `postgres_db_connection_timeouts_total` is what it turns into. The held + duration, `postgres_db_connection_held_duration_seconds`, is query latency seen from + the pool: it rises when the database slows down, whether or not the pool is under + pressure. `postgres_db_connection_checkout_duration_seconds` is the held duration under + a name that says wait; it is deprecated, and reading it as the wait is the mistake it + invites. ### Fixed since 0.2.0 -Three published entry points raise before they reach the database on 0.2.0, and on 0.2.1 -the PgBouncer-safe defaults cannot connect through PgBouncer. All work on current versions; -the workaround column is what to do on the version the row names. +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 +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. | Call | What it does on that version | Workaround there | |---|---|---| @@ -425,6 +456,8 @@ the workaround column is what to do on the version the row names. | `uow.transaction(isolation_level=…)`, `uow.managed_session(isolation_level=…)`, `uow.query(isolation_level=…)` (0.2.0) | `InvalidRequestError: This connection has already initialized a SQLAlchemy Transaction()… isolation_level may not be altered` — the level was applied after the connection had autobegun | set the level on the engine: `AsyncSessionManager(..., isolation_level="SERIALIZABLE")` or `QuerySettings(isolation_level=...)` | | `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_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`, `REPEATABLE_READ`, `SERIALIZABLE`, whose values are the PostgreSQL spellings with spaces — diff --git a/docs/guide/advanced.md b/docs/guide/advanced.md index aea1886..3f1e659 100644 --- a/docs/guide/advanced.md +++ b/docs/guide/advanced.md @@ -271,12 +271,51 @@ session_manager = create_async_session_manager( | `myapp_postgres_db_pool_size` | Gauge | Current pool size | | `myapp_postgres_db_pool_checked_out` | Gauge | Connections currently in use | | `myapp_postgres_db_pool_overflow` | Gauge | Overflow connections created | -| `myapp_postgres_db_connection_checkout_duration_seconds` | Histogram | Time to acquire connection | +| `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 | | `myapp_postgres_db_connection_errors_total` | Counter | Database errors by type (`error_type` label) | -| `myapp_postgres_db_connection_timeouts_total` | Counter | Connection timeout errors | +| `myapp_postgres_db_connection_timeouts_total` | Counter | Pool checkout timeouts, and `TimeoutError`s during execution | Without a `prefix` the names are `postgres_db_pool_size` and so on. +**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 +outage — it climbs while the pool saturates, and turns into +`connection_timeouts_total` when it crosses `pool_timeout`. The held duration is the time +between the `checkout` and `checkin` events, which is query duration seen from the pool: it +climbs when the database slows down, whether or not the pool is under any pressure. A +saturated pool moves both, and only the wait tells you which one caused the other. + +**`connection_checkout_duration_seconds` is deprecated.** It has always measured the +*held* duration, despite the name. It is still published, unchanged, so existing dashboards +keep working, and it will be removed in a future release. Point dashboards at +`connection_held_duration_seconds` for the same numbers, or at +`connection_checkout_wait_seconds` if what you actually wanted was the wait. + +The wait histogram and the pool-timeout half of the counter are not event listeners — +SQLAlchemy fires no event when a checkout is *requested* — so they come from the pool class +the session manager builds. They are populated on any engine created by +`AsyncSessionManager`, `AsyncSessionManagerBuilder` or `create_async_session_manager`. On +an engine you build yourself, wrap the pool class: + +```python +from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy_foundation_kit import resolve_pool_class +from sqlalchemy_foundation_kit.session.manager import attach_metrics, instrument_pool_class + +metrics = PostgresMetrics(prefix="myapp") +engine = create_async_engine( + url, + poolclass=instrument_pool_class(resolve_pool_class("async_adapted_queue"), metrics), +) +attach_metrics(engine, metrics) +``` + +`attach_metrics` logs a warning if the engine's pool was not built this way, rather than +leaving you with a wait histogram that never moves. + **Expose metrics endpoint:** ```python @@ -295,14 +334,21 @@ async def metrics_endpoint(): ```promql # Pool utilization -(myapp_postgres_pool_checked_out / myapp_postgres_pool_size) * 100 +(myapp_postgres_db_pool_checked_out / myapp_postgres_db_pool_size) * 100 + +# P95 wait for a connection -- the one to alert on +histogram_quantile(0.95, + rate(myapp_postgres_db_connection_checkout_wait_seconds_bucket[5m])) + +# Checkouts that gave up waiting +rate(myapp_postgres_db_connection_timeouts_total[5m]) -# P95 checkout latency -histogram_quantile(0.95, - rate(myapp_postgres_checkout_duration_seconds_bucket[5m])) +# P95 time a connection was held -- query latency, seen from the pool +histogram_quantile(0.95, + rate(myapp_postgres_db_connection_held_duration_seconds_bucket[5m])) # Error rate -rate(myapp_postgres_errors_total[5m]) +rate(myapp_postgres_db_connection_errors_total[5m]) ``` ### OpenTelemetry Tracing diff --git a/docs/guide/configuration.md b/docs/guide/configuration.md index e92988c..0df3cff 100644 --- a/docs/guide/configuration.md +++ b/docs/guide/configuration.md @@ -560,6 +560,14 @@ BasePostgresConfig( ) ``` +The series to watch is `postgres_db_connection_checkout_wait_seconds` — how long callers +wait for a connection. It climbs before anything fails, and when it reaches the `timeout` +above, callers start getting `sqlalchemy.exc.TimeoutError` and +`postgres_db_connection_timeouts_total` starts moving. A rising +`postgres_db_connection_held_duration_seconds` with a flat wait is a slower database, not a +small pool; a rising wait with a flat held duration is a pool that needs more connections. +See [Advanced → Prometheus Metrics](advanced.md#prometheus-metrics) for the full list. + ## Next Steps - **[Advanced Usage](advanced.md)** — Unit of Work, metrics, telemetry diff --git a/docs/index.md b/docs/index.md index 2f66dfd..ede9980 100644 --- a/docs/index.md +++ b/docs/index.md @@ -33,7 +33,7 @@ Every SQLAlchemy-based service typically needs: 1. **Configuration management** — DSN construction, pool settings, query options 2. **Session lifecycle** — Context managers, commit/rollback logic, cleanup 3. **Transaction management** — Unit of Work pattern with nested transactions -4. **Observability** — Metrics for pool size, checkout duration, query errors +4. **Observability** — Metrics for pool size, checkout wait, query errors 5. **Base models** — Naming conventions, timestamp mixins, custom types 6. **DI wiring** — Providers for session makers, UoW, repositories @@ -225,7 +225,7 @@ Pydantic-based configuration models: Prometheus metrics for connection pool: -- `PostgresMetrics` — Pool size, checked out connections, checkout duration, errors +- `PostgresMetrics` — Pool size, checked out connections, checkout wait, held duration, errors - Tracks health checks, pool exhaustion, timeouts #### `contrib.di` (requires `[dishka]`) diff --git a/docs/reference/index.md b/docs/reference/index.md index 866297a..bac4f10 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -140,6 +140,7 @@ Async session manager with connection pooling and health checks. members: - AsyncSessionManager - attach_metrics + - instrument_pool_class ::: sqlalchemy_foundation_kit.session.builder options: @@ -324,6 +325,7 @@ Observability protocols for monitoring. - PostgresMetricsProtocol - PoolStatsRecorder - CheckoutRecorder + - CheckoutWaitRecorder - ErrorRecorder ### Example Usage @@ -344,7 +346,14 @@ class CustomMetrics: pass def record_checkout(self, duration: float) -> None: - # Record connection checkout duration + # Record how long a connection was held, checkout to checkin + pass + + def record_checkout_wait(self, duration: float, timed_out: bool = False) -> None: + # Record how long a caller waited for a connection, and whether it gave up. + # Optional: this method is CheckoutWaitRecorder, which PostgresMetricsProtocol + # deliberately does not require. Implement it and the session manager wraps the + # pool class so the wait and the pool checkout timeout are recorded too. pass def record_error(self, error_type: str, is_timeout: bool) -> None: diff --git a/sqlalchemy_foundation_kit/__init__.py b/sqlalchemy_foundation_kit/__init__.py index 8dbe7d3..56ba530 100644 --- a/sqlalchemy_foundation_kit/__init__.py +++ b/sqlalchemy_foundation_kit/__init__.py @@ -31,6 +31,7 @@ # Protocols from .protocols import ( CheckoutRecorder, + CheckoutWaitRecorder, ErrorRecorder, PoolStatsRecorder, PostgresMetricsProtocol, @@ -89,6 +90,7 @@ "QuerySettingsProtocol", # Protocols "CheckoutRecorder", + "CheckoutWaitRecorder", "ErrorRecorder", "PoolStatsRecorder", "PostgresMetricsProtocol", diff --git a/sqlalchemy_foundation_kit/contrib/metrics/postgres.py b/sqlalchemy_foundation_kit/contrib/metrics/postgres.py index c1a4d27..95a9b8d 100644 --- a/sqlalchemy_foundation_kit/contrib/metrics/postgres.py +++ b/sqlalchemy_foundation_kit/contrib/metrics/postgres.py @@ -11,7 +11,7 @@ except ImportError: HAS_PROMETHEUS = False -# Default buckets for connection checkout duration (seconds) +# Default buckets for how long a connection was held (seconds) CONNECTION_CHECKOUT_BUCKETS: tuple[float, ...] = ( 0.001, 0.005, @@ -26,6 +26,25 @@ 5.0, ) +# Default buckets for the wait for a connection from the pool (seconds). The range runs to +# 30 s, the default ``pool_timeout``, so a saturated pool fills the top buckets instead of +# collapsing into +Inf just when the histogram is worth reading. +CONNECTION_WAIT_BUCKETS: tuple[float, ...] = ( + 0.001, + 0.005, + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + 30.0, +) + def _check_prometheus() -> None: """Check if prometheus-client is installed.""" @@ -43,12 +62,19 @@ class PostgresMetrics: - postgres_db_pool_size: Current database connection pool size. - 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_connection_checkout_duration_seconds: Time to acquire connection from pool. + - 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, + published unchanged for one more minor release. - postgres_db_connection_timeouts_total: Number of connection checkout timeouts. - postgres_db_connection_errors_total: Number of connection errors. Labels: - error_type: Type of connection error (for errors_total). + + 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``). """ def __init__(self, prefix: str | None = None) -> None: @@ -74,9 +100,19 @@ def __init__(self, prefix: str | None = None) -> None: _make_metric_name("postgres_db_pool_overflow", prefix), "Number of connections over pool_size (within max_overflow)", ) + self.connection_checkout_wait = Histogram( + _make_metric_name("postgres_db_connection_checkout_wait_seconds", prefix), + "Time a caller waited to acquire a connection from the pool", + buckets=list(CONNECTION_WAIT_BUCKETS), + ) + self.connection_held_duration = Histogram( + _make_metric_name("postgres_db_connection_held_duration_seconds", prefix), + "Time a connection was held by its caller, from checkout to checkin", + buckets=list(CONNECTION_CHECKOUT_BUCKETS), + ) self.connection_checkout_duration = Histogram( _make_metric_name("postgres_db_connection_checkout_duration_seconds", prefix), - "Time to acquire connection from pool", + "Deprecated: held duration, published under its old name. Use postgres_db_connection_held_duration_seconds", buckets=list(CONNECTION_CHECKOUT_BUCKETS), ) self.connection_timeouts_total = Counter( @@ -104,9 +140,29 @@ def record_checkout( self, duration: float, ) -> None: - """Record a database connection checkout from the pool.""" + """Record how long a connection was held, from checkout to checkin. + + Observed into ``postgres_db_connection_held_duration_seconds`` and, until the + deprecated name is dropped, into ``postgres_db_connection_checkout_duration_seconds`` + as well. + """ + self.connection_held_duration.observe(duration) self.connection_checkout_duration.observe(duration) + def record_checkout_wait( + self, + duration: float, + timed_out: bool = False, + ) -> None: + """Record the wait for a connection from the pool, and count a pool timeout. + + A pool checkout timeout is raised before any DBAPI call, so it never reaches the + engine's ``handle_error`` listener; this is the only place it is counted. + """ + self.connection_checkout_wait.observe(duration) + if timed_out: + self.connection_timeouts_total.inc() + def record_error( self, error_type: str, diff --git a/sqlalchemy_foundation_kit/protocols/__init__.py b/sqlalchemy_foundation_kit/protocols/__init__.py index a248a49..4f786cf 100644 --- a/sqlalchemy_foundation_kit/protocols/__init__.py +++ b/sqlalchemy_foundation_kit/protocols/__init__.py @@ -8,6 +8,7 @@ from .metrics import ( CheckoutRecorder, + CheckoutWaitRecorder, ErrorRecorder, PoolStatsRecorder, PostgresMetricsProtocol, @@ -15,6 +16,7 @@ __all__ = [ "CheckoutRecorder", + "CheckoutWaitRecorder", "ErrorRecorder", "PoolStatsRecorder", "PostgresMetricsProtocol", diff --git a/sqlalchemy_foundation_kit/protocols/metrics.py b/sqlalchemy_foundation_kit/protocols/metrics.py index f23f616..33467f9 100644 --- a/sqlalchemy_foundation_kit/protocols/metrics.py +++ b/sqlalchemy_foundation_kit/protocols/metrics.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import Protocol +from typing import Protocol, runtime_checkable class PoolStatsRecorder(Protocol): @@ -30,13 +30,39 @@ def record_pool_stats( class CheckoutRecorder(Protocol): - """Capability protocol for recording connection checkout duration.""" + """Capability protocol for recording how long a connection was held.""" def record_checkout(self, duration: float) -> None: - """Record a database connection checkout from the pool. + """Record a completed database connection checkout. Args: - duration: Time taken to acquire the connection from the pool, in seconds. + duration: Time between the connection leaving the pool and coming back to it, + in seconds — how long the caller *held* it, which is the query time seen + from the pool. The time a caller spent *waiting* for it is + :class:`CheckoutWaitRecorder`. + """ + ... + + +@runtime_checkable +class CheckoutWaitRecorder(Protocol): + """Capability protocol for recording the wait for a connection from the pool. + + Deliberately not part of :class:`PostgresMetricsProtocol`: a metrics object written + against the three older capabilities stays valid and simply publishes no wait series. + The session manager tests for this protocol with ``isinstance`` and instruments the + pool class only when it is satisfied, which is what ``runtime_checkable`` is for here. + """ + + def record_checkout_wait(self, duration: float, timed_out: bool = False) -> None: + """Record the time a caller spent acquiring a connection from the pool. + + Args: + duration: Seconds spent inside ``pool.connect()`` — the queue wait, and the + pre-ping round trip and connect handshake when the pool had to grow. + timed_out: True if the wait ended in ``sqlalchemy.exc.TimeoutError`` rather + than in a connection. This is the pool checkout timeout; it is raised + before any DBAPI call and so never reaches the engine's ``handle_error``. """ ... @@ -60,10 +86,15 @@ class PostgresMetricsProtocol(PoolStatsRecorder, CheckoutRecorder, ErrorRecorder Aggregates the narrow capability protocols for convenience. Implementations that only need a subset can implement the individual protocols directly. + :class:`CheckoutWaitRecorder` is intentionally left out so that implementations + written before it keep satisfying this protocol; add ``record_checkout_wait`` to get + the checkout wait histogram and the pool timeout counter as well. + Examples: >>> class MyMetrics: ... def record_pool_stats(self, pool_size: int, pool_checked_out: int, pool_overflow: int) -> None: ... ... def record_checkout(self, duration: float) -> None: ... + ... def record_checkout_wait(self, duration: float, timed_out: bool = False) -> None: ... ... def record_error(self, error_type: str, is_timeout: bool = False) -> None: ... >>> metrics: PostgresMetricsProtocol = MyMetrics() """ @@ -71,6 +102,7 @@ class PostgresMetricsProtocol(PoolStatsRecorder, CheckoutRecorder, ErrorRecorder __all__ = [ "CheckoutRecorder", + "CheckoutWaitRecorder", "ErrorRecorder", "PoolStatsRecorder", "PostgresMetricsProtocol", diff --git a/sqlalchemy_foundation_kit/session/manager.py b/sqlalchemy_foundation_kit/session/manager.py index 097008f..4c6595e 100644 --- a/sqlalchemy_foundation_kit/session/manager.py +++ b/sqlalchemy_foundation_kit/session/manager.py @@ -21,6 +21,7 @@ from .._typing import SessionT from ..base import build_engine_kwargs, resolve_pool_class +from ..protocols import CheckoutWaitRecorder if TYPE_CHECKING: from sqlalchemy.engine import Connection @@ -32,6 +33,10 @@ DEFAULT_DISPOSE_TIMEOUT_SECONDS: float = 30.0 +# Set on the pool subclass instrument_pool_class() builds, so attach_metrics() can tell an +# engine whose pool records the checkout wait from one that silently will not. +_WAIT_RECORDER_ATTRIBUTE = "_checkout_wait_recorder" + def _safe_metric_call(func: Callable[[], None], error_msg: str) -> None: """Call a metric-recording function and swallow exceptions. @@ -49,18 +54,94 @@ def _safe_metric_call(func: Callable[[], None], error_msg: str) -> None: logger.exception(error_msg) +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. + + SQLAlchemy has no event for "a checkout was requested", only for one that succeeded, so + the time a caller spends queued for a connection cannot be observed from a listener. + The returned subclass wraps :meth:`sqlalchemy.pool.Pool.connect`, which runs exactly + once per checkout — for a sync and an async engine alike, since + ``Engine.raw_connection()`` goes through it — and through which the pool's + ``TimeoutError`` passes exactly once. Under an async pool the queue wait is an + ``await_only`` inside that call, so the timer spans the whole suspension and measures + the time the caller really waited. + + ``_do_get`` is the wrong hook: ``QueuePool`` recurses into it when the non-blocking get + comes back empty and the race for an overflow slot is lost, which is exactly the + contended case worth measuring, and the wait would be observed twice. + + The recorder lives on the class, so a pool produced by ``recreate()`` after + ``dispose()`` keeps recording. + + Args: + pool_class: Pool class to instrument — anything :func:`resolve_pool_class` returns, + including a class registered with :func:`register_pool_class`. + metrics: Recorder to call once per checkout. A raising recorder is logged and + swallowed; the checkout itself is never affected. + + Returns: + A subclass of ``pool_class`` suitable for ``create_async_engine(poolclass=...)``. + + Examples: + >>> poolclass = instrument_pool_class(resolve_pool_class("async_adapted_queue"), metrics) + >>> engine = create_async_engine(url, poolclass=poolclass) + """ + + class InstrumentedPool(pool_class): # type: ignore[misc, valid-type] + """Pool that reports how long each checkout waited.""" + + def connect(self) -> Any: + started = time.perf_counter() + try: + connection = super().connect() + except SATimeoutError: + _safe_metric_call( + lambda: metrics.record_checkout_wait( + duration=time.perf_counter() - started, + timed_out=True, + ), + "Failed to record database checkout timeout", + ) + raise + _safe_metric_call( + lambda: metrics.record_checkout_wait(duration=time.perf_counter() - started), + "Failed to record database checkout wait", + ) + return connection + + setattr(InstrumentedPool, _WAIT_RECORDER_ATTRIBUTE, metrics) + InstrumentedPool.__name__ = f"Instrumented{pool_class.__name__}" + InstrumentedPool.__qualname__ = InstrumentedPool.__name__ + return InstrumentedPool + + def attach_metrics(engine: AsyncEngine, metrics: PostgresMetricsProtocol) -> None: """Attach metrics event listeners to a SQLAlchemy engine. Registers event handlers for connection checkout, checkin, and error events to collect pool statistics and connection metrics. + 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 + logged as a warning rather than left as a series that never moves. + Args: engine: SQLAlchemy ``AsyncEngine`` to attach listeners to. metrics: Metrics collector implementing ``PostgresMetricsProtocol``. """ pool = engine.pool + if isinstance(metrics, CheckoutWaitRecorder) and getattr(type(pool), _WAIT_RECORDER_ATTRIBUTE, None) is None: + logger.warning( + "%s records the connection checkout wait, but the engine's pool (%s) was not built by " + "instrument_pool_class(), so no wait and no pool checkout timeout will ever be recorded. " + "Build the engine through AsyncSessionManager, or pass " + "instrument_pool_class(pool_class, metrics) as poolclass to create_async_engine().", + type(metrics).__name__, + type(pool).__name__, + ) + def record_pool_stats() -> None: _safe_metric_call( lambda: metrics.record_pool_stats( @@ -171,7 +252,10 @@ def __init__( isolation_level: Default transaction isolation level (default: None). pool_settings: Pool configuration settings (default: None). use_orjson: If True, use orjson for JSON serialization (default: False). - metrics: Optional metrics collector (default: None). + metrics: Optional metrics collector (default: None). One that also implements + ``CheckoutWaitRecorder`` gets the pool class wrapped by + :func:`instrument_pool_class`, which is what feeds the checkout wait and + the pool checkout timeout — neither is reachable from a pool event. on_engine_created: Optional callback invoked with ``AsyncEngine`` after creation. Use for OpenTelemetry instrumentation, custom event listeners, etc. dispose_timeout: Maximum seconds to wait for engine disposal in :meth:`aclose` @@ -187,6 +271,10 @@ def __init__( self._close_lock = asyncio.Lock() self._dispose_timeout = dispose_timeout resolved_poolclass = resolve_pool_class(poolclass) + if metrics is not None and isinstance(metrics, CheckoutWaitRecorder): + # The wait for a connection is not observable from a pool event, so it has to + # be timed inside the pool -- which means deciding the class before the engine. + resolved_poolclass = instrument_pool_class(resolved_poolclass, metrics) engine_kwargs = build_engine_kwargs( echo=echo, poolclass=resolved_poolclass, diff --git a/tests/integration/test_metrics_integration.py b/tests/integration/test_metrics_integration.py new file mode 100644 index 0000000..daa7859 --- /dev/null +++ b/tests/integration/test_metrics_integration.py @@ -0,0 +1,205 @@ +"""Integration tests for the connection pool metrics, against a pool that runs out.""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import AsyncIterator +from dataclasses import dataclass + +import pytest +from sqlalchemy import text +from sqlalchemy.exc import TimeoutError as SATimeoutError +from sqlalchemy.ext.asyncio import AsyncSession + +from sqlalchemy_foundation_kit.session.manager import AsyncSessionManager +from tests.integration.conftest import PostgresContainer, asyncpg_url + +POOL_TIMEOUT_SECONDS = 0.5 +SHORT_HOLD_SECONDS = 0.25 # shorter than the timeout: whoever queues behind it is served +LONG_HOLD_SECONDS = 0.9 # longer: whoever queues behind it gives up + + +@dataclass +class _PoolSettings: + """One connection, no overflow, a short checkout timeout — a pool that runs out.""" + + kind: str = "async_adapted_queue" + size: int = 1 + max_overflow: int = 0 + pre_ping: bool = False + recycle: int = -1 + timeout: float = POOL_TIMEOUT_SECONDS + + +class _RecordingMetrics: + """Keeps every recorded value, so a test can compare the wait against the hold.""" + + def __init__(self) -> None: + self.waits: list[tuple[float, bool]] = [] + self.held: list[float] = [] + + def record_pool_stats(self, pool_size: int, pool_checked_out: int, pool_overflow: int) -> None: ... + + def record_checkout(self, duration: float) -> None: + self.held.append(duration) + + def record_checkout_wait(self, duration: float, timed_out: bool = False) -> None: + self.waits.append((duration, timed_out)) + + def record_error(self, error_type: str, is_timeout: bool = False) -> None: ... + + def reset(self) -> None: + self.waits.clear() + self.held.clear() + + @property + def served_waits(self) -> list[float]: + return [duration for duration, gave_up in self.waits if not gave_up] + + @property + def timed_out_waits(self) -> list[float]: + return [duration for duration, gave_up in self.waits if gave_up] + + +async def _holds_the_connection(manager: AsyncSessionManager[AsyncSession], seconds: float) -> None: + """Take the pool's only connection and sit on it, without asking the server to wait.""" + async with manager.get_session() as session: + await session.execute(text("SELECT 1")) + await asyncio.sleep(seconds) + + +async def _queues_behind_it(manager: AsyncSessionManager[AsyncSession]) -> None: + """Ask for a connection while the holder has it, then give it straight back.""" + await asyncio.sleep(0.05) + async with manager.get_session() as session: + await session.execute(text("SELECT 1")) + + +async def _gives_up(manager: AsyncSessionManager[AsyncSession]) -> None: + """Ask for a connection that never comes and fail the way the reporter's callers did.""" + await asyncio.sleep(0.05) + with pytest.raises(SATimeoutError): + async with manager.get_session() as session: + await session.execute(text("SELECT 1")) + + +async def _warm(manager: AsyncSessionManager[AsyncSession]) -> None: + """Pay for the connect handshake up front, so later waits are the queue wait alone.""" + async with manager.get_session() as session: + await session.execute(text("SELECT 1")) + + +@pytest.fixture +async def recording_manager( + postgres_container: PostgresContainer, +) -> AsyncIterator[tuple[AsyncSessionManager[AsyncSession], _RecordingMetrics]]: + metrics = _RecordingMetrics() + manager: AsyncSessionManager[AsyncSession] = AsyncSessionManager( + asyncpg_url(postgres_container), + poolclass="async_adapted_queue", + metrics=metrics, + pool_settings=_PoolSettings(), + ) + await _warm(manager) + metrics.reset() + yield manager, metrics + await manager.aclose() + + +# ============================================================================ +# Pool Checkout Wait and Timeout Tests +# ============================================================================ + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__pool_metrics__pool_runs_out__counts_the_timeout_and_measures_the_wait( + recording_manager: tuple[AsyncSessionManager[AsyncSession], _RecordingMetrics], +) -> None: + # Arrange: the reporter's shape with one connection instead of four — eight workers + # against a pool of four tells the same story with more noise. Two runs: one where the + # caller behind the holder is served, one where it gives up. + manager, metrics = recording_manager + + # Act: the holder keeps the only connection for less than the pool timeout. + await asyncio.gather( + _holds_the_connection(manager, SHORT_HOLD_SECONDS), + _queues_behind_it(manager), + ) + + # Assert: one caller took the connection straight away, the other sat in the queue for + # most of the first one's hold. That is a wait — and it is the mirror image of what the + # checkin listener sees, where the caller that waited longest held it for no time at + # all. Reading one off the other is the mistake the single old histogram invited. + assert len(metrics.waits) == 2, "the pool recorded no checkout wait at all" + assert metrics.timed_out_waits == [] + assert min(metrics.served_waits) < SHORT_HOLD_SECONDS / 2 + assert max(metrics.served_waits) >= SHORT_HOLD_SECONDS / 2 + assert max(metrics.held) >= SHORT_HOLD_SECONDS * 0.8 + assert min(metrics.held) < SHORT_HOLD_SECONDS / 2 + + # Arrange: now the holder keeps it for longer than the pool timeout. + metrics.reset() + + # Act + await asyncio.gather( + _holds_the_connection(manager, LONG_HOLD_SECONDS), + _gives_up(manager), + ) + + # Assert: the pool TimeoutError is raised before any DBAPI call, so it never reaches + # the engine's handle_error listener. This is the only place it is counted, and the + # counter the reporter watched stay at zero moves by exactly one. + assert len(metrics.timed_out_waits) == 1 + assert metrics.timed_out_waits[0] == pytest.approx(POOL_TIMEOUT_SECONDS, abs=0.2) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test__postgres_metrics__pool_runs_out__moves_the_series_a_dashboard_scrapes( + postgres_container: PostgresContainer, +) -> None: + # Arrange: the same run read the way an application reads it — off the Prometheus + # registry, by name. + 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=_PoolSettings(), + ) + + async with manager: + # Act + await asyncio.gather( + _holds_the_connection(manager, LONG_HOLD_SECONDS), + _gives_up(manager), + ) + + # Assert + exposed = _exposed(prefix) + assert exposed[f"{prefix}_postgres_db_connection_timeouts_total"] == 1.0 + assert exposed[f"{prefix}_postgres_db_connection_checkout_wait_seconds_count"] == 2.0 + assert exposed[f"{prefix}_postgres_db_connection_checkout_wait_seconds_sum"] >= POOL_TIMEOUT_SECONDS + assert exposed[f"{prefix}_postgres_db_connection_held_duration_seconds_sum"] >= LONG_HOLD_SECONDS * 0.8 + # The old name keeps its data for one more minor release. + assert ( + exposed[f"{prefix}_postgres_db_connection_checkout_duration_seconds_sum"] + == exposed[f"{prefix}_postgres_db_connection_held_duration_seconds_sum"] + ) + + +def _exposed(prefix: str) -> dict[str, float]: + """Read the sample values Prometheus would scrape for one metrics prefix.""" + from prometheus_client import REGISTRY, generate_latest + + samples: dict[str, float] = {} + for line in generate_latest(REGISTRY).decode().splitlines(): + if line.startswith(f"{prefix}_"): + name, value = line.rsplit(" ", 1) + samples[name] = float(value) + return samples diff --git a/tests/unit/contrib/metrics/test_postgres.py b/tests/unit/contrib/metrics/test_postgres.py index 9c0d130..8c6ec24 100644 --- a/tests/unit/contrib/metrics/test_postgres.py +++ b/tests/unit/contrib/metrics/test_postgres.py @@ -11,6 +11,7 @@ from sqlalchemy_foundation_kit.contrib.metrics.postgres import ( + CONNECTION_WAIT_BUCKETS, PostgresMetrics, _make_metric_name, ) @@ -21,6 +22,15 @@ def _unique_prefix() -> str: return f"test_{uuid.uuid4().hex[:8]}" +def _histogram(histogram: object, suffix: str) -> float: + """Read the ``_count`` or ``_sum`` sample a histogram currently exposes.""" + for metric in histogram.collect(): # type: ignore[attr-defined] + for sample in metric.samples: + if sample.name.endswith(suffix): + return float(sample.value) + raise AssertionError(f"no {suffix} sample on the histogram") + + # ============================================================================ # _make_metric_name Tests # ============================================================================ @@ -251,6 +261,114 @@ def test__postgres_metrics__record_checkout__zero_duration__succeeds() -> None: assert metrics.connection_checkout_duration is not None +# ============================================================================ +# record_checkout_wait Tests +# ============================================================================ + + +@pytest.mark.unit +def test__postgres_metrics__record_checkout_wait__observes_into_the_wait_histogram() -> None: + # Arrange + metrics = PostgresMetrics(prefix=_unique_prefix()) + + # Act + metrics.record_checkout_wait(duration=0.42) + + # Assert + assert _histogram(metrics.connection_checkout_wait, "_count") == 1.0 + assert _histogram(metrics.connection_checkout_wait, "_sum") == pytest.approx(0.42) + + +@pytest.mark.unit +def test__postgres_metrics__record_checkout_wait__timed_out__counts_the_pool_timeout() -> None: + # Arrange: the pool checkout timeout never reaches handle_error, so this is the only + # place it can be counted. + metrics = PostgresMetrics(prefix=_unique_prefix()) + + # Act + metrics.record_checkout_wait(duration=30.0, timed_out=True) + + # Assert + assert metrics.connection_timeouts_total._value.get() == 1.0 # type: ignore[attr-defined] + assert _histogram(metrics.connection_checkout_wait, "_count") == 1.0 + + +@pytest.mark.unit +def test__postgres_metrics__record_checkout_wait__served__does_not_count_a_timeout() -> None: + # Arrange + metrics = PostgresMetrics(prefix=_unique_prefix()) + + # Act + metrics.record_checkout_wait(duration=0.01) + + # Assert + assert metrics.connection_timeouts_total._value.get() == 0.0 # type: ignore[attr-defined] + + +@pytest.mark.unit +def test__postgres_metrics__wait_and_held__are_separate_series() -> None: + # Arrange: a caller that waited 2 s and then held the connection for 0.05 s. Reading + # the wait off the held histogram is exactly the mistake the old single series invited. + metrics = PostgresMetrics(prefix=_unique_prefix()) + + # Act + metrics.record_checkout_wait(duration=2.0) + metrics.record_checkout(duration=0.05) + + # Assert + assert _histogram(metrics.connection_checkout_wait, "_sum") == pytest.approx(2.0) + assert _histogram(metrics.connection_held_duration, "_sum") == pytest.approx(0.05) + + +@pytest.mark.unit +def test__postgres_metrics__record_checkout__also_feeds_the_deprecated_name() -> None: + # Arrange + metrics = PostgresMetrics(prefix=_unique_prefix()) + + # Act + metrics.record_checkout(duration=0.25) + + # Assert + assert _histogram(metrics.connection_held_duration, "_sum") == pytest.approx(0.25) + assert _histogram(metrics.connection_checkout_duration, "_sum") == pytest.approx(0.25) + + +@pytest.mark.unit +def test__postgres_metrics__wait_buckets__reach_the_default_pool_timeout() -> None: + # Arrange & Act & Assert: a saturated pool times out at pool_timeout, 30 s by default. + # A wait histogram that stopped at 5 s would put every one of them in +Inf. + assert CONNECTION_WAIT_BUCKETS[-1] >= 30.0 + assert tuple(sorted(CONNECTION_WAIT_BUCKETS)) == CONNECTION_WAIT_BUCKETS + + +@pytest.mark.unit +def test__postgres_metrics__publishes_the_documented_series() -> None: + # Arrange + from prometheus_client import REGISTRY, generate_latest + + prefix = _unique_prefix() + + # Act + PostgresMetrics(prefix=prefix) + published = { + line.split(" ")[2] + for line in generate_latest(REGISTRY).decode().splitlines() + if line.startswith(f"# HELP {prefix}_") and not line.split(" ")[2].endswith("_created") + } + + # Assert + assert published == { + f"{prefix}_postgres_db_pool_size", + f"{prefix}_postgres_db_pool_checked_out", + f"{prefix}_postgres_db_pool_overflow", + f"{prefix}_postgres_db_connection_checkout_wait_seconds", + f"{prefix}_postgres_db_connection_held_duration_seconds", + f"{prefix}_postgres_db_connection_checkout_duration_seconds", + f"{prefix}_postgres_db_connection_timeouts_total", + f"{prefix}_postgres_db_connection_errors_total", + } + + # ============================================================================ # record_error Tests # ============================================================================ diff --git a/tests/unit/session/test_manager.py b/tests/unit/session/test_manager.py index a0760d8..bffed54 100644 --- a/tests/unit/session/test_manager.py +++ b/tests/unit/session/test_manager.py @@ -2,11 +2,15 @@ from __future__ import annotations +import threading +import time +from typing import cast from unittest.mock import AsyncMock, Mock, patch import pytest from sqlalchemy.exc import TimeoutError as SATimeoutError from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.pool import QueuePool from sqlalchemy_foundation_kit.session.manager import ( DEFAULT_DISPOSE_TIMEOUT_SECONDS, @@ -14,6 +18,7 @@ _safe_metric_call, attach_metrics, attach_search_path, + instrument_pool_class, ) # ============================================================================ @@ -295,6 +300,189 @@ def test__attach_metrics__pool_without_size__uses_zero() -> None: assert call_args["pool_overflow"] == 0 +# ============================================================================ +# instrument_pool_class Tests +# ============================================================================ + + +class _WaitRecorder: + """Records every checkout wait the pool reports, in order.""" + + def __init__(self) -> None: + self.waits: list[tuple[float, bool]] = [] + + def record_checkout_wait(self, duration: float, timed_out: bool = False) -> None: + self.waits.append((duration, timed_out)) + + @property + def timeouts(self) -> int: + return sum(1 for _duration, timed_out in self.waits if timed_out) + + +def _pool(recorder: _WaitRecorder, timeout: float = 0.1) -> QueuePool: + """A one-connection QueuePool over a fake DBAPI, instrumented with ``recorder``.""" + poolclass = instrument_pool_class(QueuePool, recorder) + return cast( + "QueuePool", + poolclass(Mock, pool_size=1, max_overflow=0, timeout=timeout), + ) + + +def test__instrument_pool_class__returns_subclass_of_the_original() -> None: + # Arrange & Act + poolclass = instrument_pool_class(QueuePool, _WaitRecorder()) + + # Assert + assert issubclass(poolclass, QueuePool) + assert poolclass is not QueuePool + assert poolclass.__name__ == "InstrumentedQueuePool" + + +def test__instrument_pool_class__successful_checkout__records_the_wait() -> None: + # Arrange + recorder = _WaitRecorder() + pool = _pool(recorder) + + # Act + connection = pool.connect() + connection.close() + + # Assert + assert len(recorder.waits) == 1 + duration, timed_out = recorder.waits[0] + assert timed_out is False + assert duration >= 0.0 + + +def test__instrument_pool_class__exhausted_pool__counts_the_timeout_and_reraises() -> None: + # Arrange: one connection, already checked out, so the next caller can only time out. + recorder = _WaitRecorder() + pool = _pool(recorder) + held = pool.connect() + + # Act + with pytest.raises(SATimeoutError): + pool.connect() + + # Assert + assert recorder.timeouts == 1 + timed_out_wait = next(duration for duration, timed_out in recorder.waits if timed_out) + assert timed_out_wait >= 0.1 + held.close() + + +def test__instrument_pool_class__measures_the_wait_not_the_time_held() -> None: + # Arrange: the first caller holds its connection for HOLD seconds; the second one is + # queued behind it the whole time and then gives the connection straight back. What + # the checkin listener sees is the mirror image of what the pool made each of them + # wait, which is the whole reason the wait needs its own metric. + hold_seconds = 0.3 + recorder = _WaitRecorder() + pool = _pool(recorder, timeout=5.0) + first = pool.connect() + + def second_caller() -> None: + pool.connect().close() + + # Act + waiter = threading.Thread(target=second_caller) + waiter.start() + time.sleep(hold_seconds) + first.close() + waiter.join() + + # Assert + assert recorder.timeouts == 0 + first_wait, second_wait = (duration for duration, _timed_out in recorder.waits) + assert first_wait < hold_seconds / 2 + assert second_wait >= hold_seconds * 0.8 + + +def test__instrument_pool_class__pool_recreated__keeps_recording() -> None: + # Arrange: dispose() replaces the pool through recreate(), which a per-instance + # recorder would not survive. + recorder = _WaitRecorder() + pool = _pool(recorder) + + # Act + recreated = pool.recreate() + recreated.connect().close() + + # Assert + assert type(recreated) is type(pool) + assert len(recorder.waits) == 1 + + +def test__instrument_pool_class__recorder_raises__checkout_still_succeeds() -> None: + # Arrange + recorder = Mock() + recorder.record_checkout_wait.side_effect = RuntimeError("metrics backend is down") + poolclass = instrument_pool_class(QueuePool, recorder) + pool = poolclass(Mock, pool_size=1, max_overflow=0) + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.logger") as mock_logger: + connection = pool.connect() + + # Assert + assert connection is not None + mock_logger.exception.assert_called_once() + connection.close() + + +def test__attach_metrics__pool_not_instrumented__warns_that_the_wait_is_lost() -> None: + # Arrange: an engine somebody built themselves, so the pool class never went through + # instrument_pool_class and the wait series would stay empty with nobody the wiser. + mock_engine = Mock() + mock_engine.pool = QueuePool(Mock) + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event"): + with patch("sqlalchemy_foundation_kit.session.manager.logger") as mock_logger: + attach_metrics(mock_engine, _WaitRecorder()) + + # Assert + mock_logger.warning.assert_called_once() + assert "instrument_pool_class" in mock_logger.warning.call_args[0][0] + + +def test__attach_metrics__instrumented_pool__does_not_warn() -> None: + # Arrange + metrics = _WaitRecorder() + mock_engine = Mock() + mock_engine.pool = instrument_pool_class(QueuePool, metrics)(Mock) + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event"): + with patch("sqlalchemy_foundation_kit.session.manager.logger") as mock_logger: + attach_metrics(mock_engine, metrics) + + # Assert + mock_logger.warning.assert_not_called() + + +def test__attach_metrics__metrics_without_wait_support__does_not_warn() -> None: + # Arrange: an implementation written against PostgresMetricsProtocol alone stays valid + # and simply publishes no wait series -- that is not worth a warning on every startup. + class OlderMetrics: + def record_pool_stats(self, pool_size: int, pool_checked_out: int, pool_overflow: int) -> None: ... + + def record_checkout(self, duration: float) -> None: ... + + def record_error(self, error_type: str, is_timeout: bool = False) -> None: ... + + mock_engine = Mock() + mock_engine.pool = QueuePool(Mock) + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.event"): + with patch("sqlalchemy_foundation_kit.session.manager.logger") as mock_logger: + attach_metrics(mock_engine, OlderMetrics()) + + # Assert + mock_logger.warning.assert_not_called() + + # ============================================================================ # attach_search_path Tests # ============================================================================ @@ -402,6 +590,48 @@ def test__async_session_manager__init__no_metrics__does_not_attach() -> None: mock_attach.assert_not_called() +def test__async_session_manager__init__metrics_record_the_wait__instruments_the_pool_class() -> None: + # Arrange + metrics = _WaitRecorder() + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.create_async_engine") as mock_create: + with patch("sqlalchemy_foundation_kit.session.manager.attach_metrics"): + AsyncSessionManager("postgresql+asyncpg://localhost/test", poolclass="queue", metrics=metrics) + + # Assert + poolclass = mock_create.call_args[1]["poolclass"] + assert issubclass(poolclass, QueuePool) + assert poolclass is not QueuePool + + +def test__async_session_manager__init__metrics_without_wait_support__leaves_the_pool_class_alone() -> None: + # Arrange + class OlderMetrics: + def record_pool_stats(self, pool_size: int, pool_checked_out: int, pool_overflow: int) -> None: ... + + def record_checkout(self, duration: float) -> None: ... + + def record_error(self, error_type: str, is_timeout: bool = False) -> None: ... + + # Act + with patch("sqlalchemy_foundation_kit.session.manager.create_async_engine") as mock_create: + with patch("sqlalchemy_foundation_kit.session.manager.attach_metrics"): + AsyncSessionManager("postgresql+asyncpg://localhost/test", poolclass="queue", metrics=OlderMetrics()) + + # Assert + assert mock_create.call_args[1]["poolclass"] is QueuePool + + +def test__async_session_manager__init__no_metrics__leaves_the_pool_class_alone() -> None: + # Arrange & Act + with patch("sqlalchemy_foundation_kit.session.manager.create_async_engine") as mock_create: + AsyncSessionManager("postgresql+asyncpg://localhost/test", poolclass="queue") + + # Assert + assert mock_create.call_args[1]["poolclass"] is QueuePool + + def test__async_session_manager__init__search_path__attaches_it() -> None: # Arrange & Act with patch("sqlalchemy_foundation_kit.session.manager.create_async_engine"):