diff --git a/src/db.py b/src/db.py index 005d52eb..633f6f06 100644 --- a/src/db.py +++ b/src/db.py @@ -12,7 +12,11 @@ from sqlalchemy.orm import declarative_base from sqlalchemy.pool import NullPool, QueuePool from src.config import settings -from src.telemetry.prometheus.metrics import db_queries_in_flight_gauge +from src.telemetry.prometheus.metrics import ( + db_connections_established_counter, + db_connections_open_gauge, + db_queries_in_flight_gauge, +) logger = logging.getLogger(__name__) @@ -220,6 +224,81 @@ def register_db_query_instrumentation(instance_type: str) -> None: _db_query_instrumentation_registered = True +class DBConnectionTracker: + """Tracks physical DB connections open on this engine via pool lifecycle events. + + Drift-proof, mirroring ``DBQueryInflightTracker``: marks the ``ConnectionRecord`` + on ``connect`` and decrements only if that mark is still present on + ``close``/``invalidate``, so each physical connection increments the gauge + exactly once and decrements at most once — it can't leak upward or go negative + when both events fire during invalidation cleanup. Works for every pool class, + including ``NullPool`` (whose pool keeps no records, so the scrape-time + ``db_pool_connections`` collector reads zero). + """ + + # Marker on ConnectionRecord.info recording that we incremented for this + # connection, so we decrement exactly once across close/invalidate. + OPEN_KEY: str = "_honcho_conn_open" + + def __init__(self, open_child: Any, established_child: Any) -> None: + self._open: Any = open_child + self._established: Any = established_child + + def on_connect(self, _dbapi_connection: Any, connection_record: Any) -> None: + try: + connection_record.info[self.OPEN_KEY] = True + self._established.inc() + self._open.inc() + except Exception: + logger.debug("db-connection gauge inc failed", exc_info=True) + + def on_close(self, _dbapi_connection: Any, connection_record: Any, *_: Any) -> None: + try: + if connection_record is not None and connection_record.info.pop( + self.OPEN_KEY, False + ): + self._open.dec() + except Exception: + logger.debug("db-connection gauge dec failed", exc_info=True) + + +# Process-wide tracker, created at registration (None until then / if metrics off). +_connection_tracker: DBConnectionTracker | None = None + + +_db_connection_instrumentation_registered = False + + +def register_db_connection_instrumentation(instance_type: str) -> None: + """Attach physical-connection tracking to the engine (no-op if metrics off). + + Counts connections via pool lifecycle events, so it reports real numbers under + any pool class — unlike the pool-object collector, which reads zero under + ``NullPool``. Pre-resolving the labeled children materializes both series at 0, + so an absent series signals a broken scrape rather than "no connections" (the + zero-init convention). Idempotent: repeated calls won't attach duplicate + listeners, which would double-count connections. + """ + global _connection_tracker, _db_connection_instrumentation_registered + if not settings.METRICS.ENABLED or _db_connection_instrumentation_registered: + return + open_child = db_connections_open_gauge.labels(instance_type=instance_type) + established_child = db_connections_established_counter.labels( + instance_type=instance_type + ) + _connection_tracker = DBConnectionTracker(open_child, established_child) + sync_engine = engine.sync_engine + event.listen(sync_engine, "connect", _connection_tracker.on_connect) + # A connection is torn down by close (normal return / recycle discard), + # invalidate (broken connection), or detach — the last fires on GC-cleanup of an + # abandoned async connection, where NullPool's close is a no-op so `close` never + # fires. All three carry the ConnectionRecord; the marker dedupes if more than + # one fires for the same connection. + for teardown_event in ("close", "invalidate", "detach"): + event.listen(sync_engine, teardown_event, _connection_tracker.on_close) + _db_connection_instrumentation_registered = True + + # Define your naming convention convention = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", # Index - supports multi-column diff --git a/src/deriver/__main__.py b/src/deriver/__main__.py index ce3969b5..9f5e4d91 100644 --- a/src/deriver/__main__.py +++ b/src/deriver/__main__.py @@ -6,7 +6,11 @@ import uvloop from prometheus_client import start_http_server from src.config import settings -from src.db import engine, register_db_query_instrumentation +from src.db import ( + engine, + register_db_connection_instrumentation, + register_db_query_instrumentation, +) from src.startup import validate_embedding_schema from src.telemetry import ( initialize_telemetry_async, @@ -26,6 +30,7 @@ def start_metrics_server() -> None: # Expose DB connection-pool stats for this deriver instance. register_db_pool_collector("deriver") register_db_query_instrumentation("deriver") + register_db_connection_instrumentation("deriver") # region ai # Zero-init bounded-label counters so a missing series signals a broken scrape, diff --git a/src/main.py b/src/main.py index 6930d4c1..a1ec9765 100644 --- a/src/main.py +++ b/src/main.py @@ -17,7 +17,12 @@ from sentry_sdk.integrations.starlette import StarletteIntegration from src._version import HONCHO_VERSION from src.cache.client import close_cache, init_cache from src.config import settings -from src.db import engine, register_db_query_instrumentation, request_context +from src.db import ( + engine, + register_db_connection_instrumentation, + register_db_query_instrumentation, + request_context, +) from src.exceptions import HonchoException from src.routers import ( conclusions, @@ -109,6 +114,7 @@ async def lifespan(_: FastAPI): # Expose DB connection-pool stats for this API instance (no-op if metrics off) register_db_pool_collector("api") register_db_query_instrumentation("api") + register_db_connection_instrumentation("api") # region ai # Zero-init bounded-label counters so a missing series signals a broken scrape, diff --git a/src/telemetry/prometheus/metrics.py b/src/telemetry/prometheus/metrics.py index 749591d9..cead893a 100644 --- a/src/telemetry/prometheus/metrics.py +++ b/src/telemetry/prometheus/metrics.py @@ -208,6 +208,27 @@ db_queries_in_flight_gauge = NamespacedGauge( ["namespace", "instance_type"], ) +# Physical DB connections, tracked via SQLAlchemy connection-lifecycle events +# (see DBConnectionTracker in src/db.py) rather than the pool object, so they are +# visible under EVERY pool class — including NullPool, whose pool holds no records +# for the scrape-time db_pool_connections collector to read. Under QueuePool this +# roughly equals db_pool_connections{checked_in} + {checked_out}; its unique value +# is under NullPool, where that collector reads zero. +db_connections_open_gauge = NamespacedGauge( + "db_connections_open", + "Physical DB connections currently open by this instance, across all pool " + + "classes (tracks concurrency of DB work under NullPool, pool occupancy " + + "under QueuePool)", + ["namespace", "instance_type"], +) + +db_connections_established_counter = NamespacedCounter( + "db_connections_established", + "Physical DB connections established since process start. Under NullPool, " + + "rate() approximates request rate (one connect per DB checkout)", + ["namespace", "instance_type"], +) + @final class PrometheusMetrics: @@ -409,6 +430,11 @@ class PrometheusMetrics: """Pre-create bounded-label counter children at 0 for this process, so an absent series means a broken scrape rather than "nothing happened". + Note: the DB-instrumentation metrics (db_queries_in_flight, + db_connections_open, db_connections_established) are NOT initialized here — + they zero-init via the pre-resolved labeled children in their register_db_* + functions in src/db.py, so an auditor should not read them as forgotten. + Args: instance_type: "api" or "deriver" — selects the process-specific counters. Event-type and buffer metrics are initialized in both. diff --git a/tests/telemetry/test_db_connection_metrics.py b/tests/telemetry/test_db_connection_metrics.py new file mode 100644 index 00000000..c9f4cfb8 --- /dev/null +++ b/tests/telemetry/test_db_connection_metrics.py @@ -0,0 +1,94 @@ +"""Tests for the physical-DB-connection metrics. + +``db_connections_open`` (gauge) and ``db_connections_established`` (counter) are +driven by SQLAlchemy connection-lifecycle events rather than the pool object, so +they report real connections under EVERY pool class — including ``NullPool``, whose +pool keeps no records for the scrape-time ``db_pool_connections`` collector to read. + +Asserts two properties: +- zero-init — resolving a labeled child materializes it at 0, so an absent series + means a broken scrape rather than "no connections"; +- ``DBConnectionTracker`` semantics — increment once per connect, decrement at most + once per connection (marker-guarded, so it can't leak upward or go negative), and + the establishment counter is monotonic (closes never decrement it). +""" + +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from prometheus_client import REGISTRY + +from src.db import DBConnectionTracker +from src.telemetry.prometheus.metrics import ( + db_connections_established_counter, + db_connections_open_gauge, +) + + +@pytest.fixture +def ns(monkeypatch: pytest.MonkeyPatch) -> str: + """Enable metrics under a namespace unique to this test. + + The process-global REGISTRY keeps a materialized child for the rest of the + session, so a shared namespace would let one test satisfy another's + presence/absence assertions independently of the code under test. + """ + namespace = f"test_db_conn_{uuid4().hex[:8]}" + monkeypatch.setattr("src.config.settings.METRICS.ENABLED", True) + monkeypatch.setattr("src.config.settings.METRICS.NAMESPACE", namespace) + return namespace + + +def sample(name: str, namespace: str, **labels: str) -> float | None: + """Value of a series if it exists, else None. Never materializes it.""" + return REGISTRY.get_sample_value(name, {"namespace": namespace, **labels}) + + +def test_connection_children_zero_init(ns: str) -> None: + """Resolving the labeled children materializes both series at 0.""" + db_connections_open_gauge.labels(instance_type="api") + db_connections_established_counter.labels(instance_type="api") + + assert sample("db_connections_open", ns, instance_type="api") == 0.0 + # prometheus_client appends _total to counter names + assert sample("db_connections_established_total", ns, instance_type="api") == 0.0 + + +def test_tracker_inc_dec_and_counter_monotonic(ns: str) -> None: + """connect increments both metrics; close decrements only the gauge.""" + open_child = db_connections_open_gauge.labels(instance_type="api") + established_child = db_connections_established_counter.labels(instance_type="api") + tracker = DBConnectionTracker(open_child, established_child) + + rec1, rec2 = SimpleNamespace(info={}), SimpleNamespace(info={}) + tracker.on_connect(None, rec1) + tracker.on_connect(None, rec2) + assert sample("db_connections_open", ns, instance_type="api") == 2.0 + assert sample("db_connections_established_total", ns, instance_type="api") == 2.0 + + tracker.on_close(None, rec1) + tracker.on_close(None, rec2) + assert sample("db_connections_open", ns, instance_type="api") == 0.0 + # the counter is monotonic: closes never decrement it + assert sample("db_connections_established_total", ns, instance_type="api") == 2.0 + + +def test_marker_prevents_double_dec_and_negative(ns: str) -> None: + """The ConnectionRecord marker bounds each connection to one dec.""" + open_child = db_connections_open_gauge.labels(instance_type="api") + established_child = db_connections_established_counter.labels(instance_type="api") + tracker = DBConnectionTracker(open_child, established_child) + + # a close with no matching connect must not drive the gauge negative + tracker.on_close(None, SimpleNamespace(info={})) + assert sample("db_connections_open", ns, instance_type="api") == 0.0 + + # connect, then close AND invalidate on the same record (both fire during + # invalidation cleanup): the marker ensures exactly one decrement. The third + # positional arg is invalidate's exception, absorbed by on_close's *_. + rec = SimpleNamespace(info={}) + tracker.on_connect(None, rec) + tracker.on_close(None, rec) + tracker.on_close(None, rec, ValueError("invalidated")) + assert sample("db_connections_open", ns, instance_type="api") == 0.0