Merge pull request #1055 from plastic-labs/phil/db-connection-metrics

feat(telemetry): add physical DB-connection metrics (visible under NullPool)

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

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,

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,

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.

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