feat(db): add acquisition counter and in-flight query gauge

Build on the pool-connection metrics with two signals that turn detection
into diagnosis under transaction-pooler saturation:

- db_connection_acquisitions{outcome=ok|retried|exhausted}: counts how often
  connection checkout retries through pooler rejection — the alertable early
  warning before requests start failing.
- db_queries_in_flight: statements actually executing on the wire (via
  SQLAlchemy cursor-execute events, drift-proof across query errors). Pairs
  with checked_out: the gap reveals connections held but parked (the "idle in
  transaction during an external call" antipattern). Labeled namespace +
  instance_type only; gated on METRICS.ENABLED for zero overhead when off.

Add DB-free unit tests for retry outcomes, polling backoff, and in-flight
gauge drift handling.

src/db.py

@@ -1,7 +1,9 @@
 import contextvars
+import logging
+from typing import Any
 
 import sentry_sdk
-from sqlalchemy import MetaData, text
+from sqlalchemy import MetaData, event, text
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.exc import TimeoutError as SQLAlchemyTimeoutError
 from sqlalchemy.ext.asyncio import (
@@ -19,6 +21,12 @@ from tenacity import (
 )
 
 from src.config import settings
+from src.telemetry.prometheus.metrics import (
+    db_queries_in_flight_gauge,
+    prometheus_metrics,
+)
+
+logger = logging.getLogger(__name__)
 
 connect_args = {"prepare_threshold": None}
 
@@ -63,6 +71,18 @@ SessionLocal = async_sessionmaker(
 # pooler surfaces "too many clients" / connection refusals).
 RETRYABLE_DB_CONNECTION_ERRORS = (SQLAlchemyTimeoutError, OperationalError)
 
+# Identifies this process ("api" | "deriver") on DB metrics. Set once at startup
+# by register_db_query_instrumentation; stays "unknown" if metrics are disabled.
+_db_instance_type: str = "unknown"
+
+
+def _record_acquisition_outcome(outcome: str) -> None:
+    """Record a connection-acquisition outcome (no-op when metrics disabled)."""
+    if settings.METRICS.ENABLED:
+        prometheus_metrics.record_db_connection_acquisition(
+            instance_type=_db_instance_type, outcome=outcome
+        )
+
 
 def get_pool_stats() -> dict[str, int]:
     """Return live connection-pool stats for this process.
@@ -106,6 +126,7 @@ async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None:
         if not settings.DB.CONNECTION_RETRY_ENABLED:
             await db.connection()
             return
+        attempts = 0
         try:
             async for attempt in AsyncRetrying(
                 wait=wait_exponential_jitter(
@@ -117,12 +138,78 @@ async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None:
                 reraise=True,
             ):
                 with attempt:
+                    attempts += 1
                     await db.connection()
         except RETRYABLE_DB_CONNECTION_ERRORS as e:
+            _record_acquisition_outcome("exhausted")
             if settings.SENTRY.ENABLED:
                 sentry_sdk.set_context("db_pool", get_pool_stats())
                 sentry_sdk.capture_exception(e)
             raise
+        # "ok" on first try, "retried" if backoff was needed before success.
+        _record_acquisition_outcome("ok" if attempts <= 1 else "retried")
+
+
+class DBQueryInflightTracker:
+    """Tracks statements executing on the wire via SQLAlchemy cursor events.
+
+    Drift-proof: marks ``Connection.info`` when a statement starts and clears it
+    on completion OR error, so the gauge can't leak upward (an errored statement
+    skips ``after_cursor_execute``) or go negative (a connect-time error has no
+    matching start). Bound to a pre-resolved labeled gauge child so the
+    per-statement hot path does no label resolution.
+    """
+
+    # Marker on Connection.info recording that we incremented for the current
+    # statement, so we decrement exactly once on completion or error.
+    INFLIGHT_KEY: str = "_honcho_inflight"
+
+    def __init__(self, gauge_child: Any) -> None:
+        self._child: Any = gauge_child
+
+    def on_before(self, conn: Any, *_: Any) -> None:
+        try:
+            conn.info[self.INFLIGHT_KEY] = True
+            self._child.inc()
+        except Exception:
+            logger.debug("in-flight gauge inc failed", exc_info=True)
+
+    def on_after(self, conn: Any, *_: Any) -> None:
+        try:
+            if conn.info.pop(self.INFLIGHT_KEY, False):
+                self._child.dec()
+        except Exception:
+            logger.debug("in-flight gauge dec failed", exc_info=True)
+
+    def on_error(self, exception_context: Any) -> None:
+        try:
+            conn = exception_context.connection
+            if conn is not None and conn.info.pop(self.INFLIGHT_KEY, False):
+                self._child.dec()
+        except Exception:
+            logger.debug("in-flight gauge error-path dec failed", exc_info=True)
+
+
+# Process-wide tracker, created at registration (None until then / if metrics off).
+_inflight_tracker: DBQueryInflightTracker | None = None
+
+
+def register_db_query_instrumentation(instance_type: str) -> None:
+    """Attach per-statement in-flight tracking to the engine (no-op if off).
+
+    Gated on METRICS.ENABLED so there is zero overhead — not even attached event
+    listeners — when metrics are disabled. Safe to call once per process.
+    """
+    global _db_instance_type, _inflight_tracker
+    _db_instance_type = instance_type
+    if not settings.METRICS.ENABLED:
+        return
+    child = db_queries_in_flight_gauge.labels(instance_type=instance_type)
+    _inflight_tracker = DBQueryInflightTracker(child)
+    sync_engine = engine.sync_engine
+    event.listen(sync_engine, "before_cursor_execute", _inflight_tracker.on_before)
+    event.listen(sync_engine, "after_cursor_execute", _inflight_tracker.on_after)
+    event.listen(sync_engine, "handle_error", _inflight_tracker.on_error)
 
 
 # Define your naming convention

src/deriver/__main__.py

@@ -6,7 +6,7 @@ import uvloop
 from prometheus_client import start_http_server
 
 from src.config import settings
-from src.db import engine
+from src.db import engine, register_db_query_instrumentation
 from src.startup import validate_embedding_schema
 from src.telemetry import (
     initialize_telemetry_async,
@@ -24,6 +24,7 @@ def start_metrics_server() -> None:
     start_http_server(9090)
     # Expose DB connection-pool stats for this deriver instance.
     register_db_pool_collector("deriver")
+    register_db_query_instrumentation("deriver")
     logger.info("Prometheus metrics server started on port 9090")
 
 

src/main.py

@@ -19,7 +19,7 @@ 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, request_context
+from src.db import engine, register_db_query_instrumentation, request_context
 from src.exceptions import HonchoException
 from src.routers import (
     conclusions,
@@ -134,6 +134,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")
 
     # Validate embedding schema before serving any traffic. Fails closed: if
     # the configured EMBEDDING_VECTOR_DIMENSIONS does not match the physical

src/telemetry/prometheus/metrics.py

@@ -127,6 +127,23 @@ telemetry_buffer_size_gauge = NamespacedGauge(
     ["namespace"],
 )
 
+# DB connection-pool health. The acquisitions counter measures how often we hit
+# (and retry through) transaction-pooler saturation; the in-flight gauge counts
+# statements actually executing on the wire, so checked_out minus in_flight
+# reveals connections held but parked (the "idle in transaction during an
+# external call" antipattern).
+db_connection_acquisitions_counter = NamespacedCounter(
+    "db_connection_acquisitions",
+    "DB connection acquisitions by outcome (ok=first try, retried, exhausted)",
+    ["namespace", "instance_type", "outcome"],
+)
+
+db_queries_in_flight_gauge = NamespacedGauge(
+    "db_queries_in_flight",
+    "DB statements currently executing on a connection for this instance",
+    ["namespace", "instance_type"],
+)
+
 
 @final
 class PrometheusMetrics:
@@ -277,6 +294,18 @@ class PrometheusMetrics:
         except Exception as e:
             self._handle_metric_error("set_telemetry_buffer_size", e)
 
+    def record_db_connection_acquisition(
+        self, *, instance_type: str, outcome: str
+    ) -> None:
+        # outcome is one of "ok" | "retried" | "exhausted".
+        try:
+            db_connection_acquisitions_counter.labels(
+                instance_type=instance_type,
+                outcome=outcome,
+            ).inc()
+        except Exception as e:
+            self._handle_metric_error("record_db_connection_acquisition", e)
+
 
 prometheus_metrics = PrometheusMetrics()
 

tests/test_db_resilience.py

@@ -0,0 +1,159 @@
+"""Unit tests for DB connection resilience + observability.
+
+These are DB-free: they exercise the retry helper against a fake session, the
+deriver polling backoff math, and the in-flight gauge listeners directly.
+"""
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+import src.db as db_module
+from src.config import settings
+from src.db import DBQueryInflightTracker, acquire_connection_with_retry
+from src.telemetry.prometheus.metrics import (
+    db_connection_acquisitions_counter,
+    db_queries_in_flight_gauge,
+)
+
+
+def _make_operational_error() -> OperationalError:
+    """A stand-in for how a saturated pooler surfaces ('too many clients')."""
+    return OperationalError("SELECT 1", {}, Exception("too many clients"))
+
+
+class _FlakyConnSession:
+    """Fake AsyncSession whose connection() fails N times then succeeds."""
+
+    def __init__(self, fail_times: int, *, always_fail: bool = False) -> None:
+        self.fail_times: int = fail_times
+        self.always_fail: bool = always_fail
+        self.calls: int = 0
+
+    async def connection(self) -> None:
+        self.calls += 1
+        if self.always_fail or self.calls <= self.fail_times:
+            raise _make_operational_error()
+
+
+def _acq_count(outcome: str) -> float:
+    child = db_connection_acquisitions_counter.labels(
+        instance_type="api", outcome=outcome
+    )
+    return float(child._value.get())  # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType]
+
+
+@pytest.fixture
+def metrics_on(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(settings.METRICS, "ENABLED", True)
+    monkeypatch.setattr(settings.METRICS, "NAMESPACE", "test")
+    monkeypatch.setattr(db_module, "_db_instance_type", "api")
+    # Fast, deterministic backoff so retry tests don't sleep.
+    monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_ENABLED", True)
+    monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS", 0.001)
+    monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_BACKOFF_MAX_SECONDS", 0.01)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("metrics_on")
+async def test_acquire_succeeds_first_try_records_ok() -> None:
+    before = _acq_count("ok")
+    session = _FlakyConnSession(fail_times=0)
+    await acquire_connection_with_retry(session, "request:test")  # pyright: ignore[reportArgumentType]
+    assert session.calls == 1
+    assert _acq_count("ok") == before + 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("metrics_on")
+async def test_acquire_retries_then_succeeds_records_retried() -> None:
+    before = _acq_count("retried")
+    session = _FlakyConnSession(fail_times=2)
+    await acquire_connection_with_retry(session, "request:test")  # pyright: ignore[reportArgumentType]
+    assert session.calls == 3  # two failures then success
+    assert _acq_count("retried") == before + 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("metrics_on")
+async def test_acquire_exhausts_budget_reraises_and_records(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # Tiny budget so the loop gives up quickly under sustained failure.
+    monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_MAX_DELAY_SECONDS", 0.05)
+    before = _acq_count("exhausted")
+    session = _FlakyConnSession(fail_times=0, always_fail=True)
+    with pytest.raises(OperationalError):
+        await acquire_connection_with_retry(session, "request:test")  # pyright: ignore[reportArgumentType]
+    assert session.calls >= 1
+    assert _acq_count("exhausted") == before + 1
+
+
+@pytest.mark.asyncio
+async def test_acquire_disabled_calls_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_ENABLED", False)
+    session = _FlakyConnSession(fail_times=0)
+    await acquire_connection_with_retry(session, "request:test")  # pyright: ignore[reportArgumentType]
+    assert session.calls == 1
+
+
+def test_polling_backoff_sequence_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True)
+    monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0)
+    monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0)
+    monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 30.0)
+
+    from src.deriver.queue_manager import QueueManager
+
+    qm = QueueManager()
+    seq = [qm._advance_poll_interval() for _ in range(8)]  # pyright: ignore[reportPrivateUsage]
+    assert seq == [1.0, 2.0, 4.0, 8.0, 16.0, 30.0, 30.0, 30.0]  # caps at max
+
+    qm._reset_poll_interval()  # pyright: ignore[reportPrivateUsage]
+    assert qm._advance_poll_interval() == 1.0  # pyright: ignore[reportPrivateUsage]
+
+
+def test_polling_backoff_disabled_stays_constant(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", False)
+    monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0)
+
+    from src.deriver.queue_manager import QueueManager
+
+    qm = QueueManager()
+    assert [qm._advance_poll_interval() for _ in range(3)] == [1.0, 1.0, 1.0]  # pyright: ignore[reportPrivateUsage]
+
+
+def test_inflight_gauge_no_drift(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(settings.METRICS, "NAMESPACE", "test")
+    child: Any = db_queries_in_flight_gauge.labels(instance_type="api")
+    tracker = DBQueryInflightTracker(child)
+    key = DBQueryInflightTracker.INFLIGHT_KEY
+
+    def value() -> float:
+        return float(child._value.get())
+
+    start = value()
+    conn = SimpleNamespace(info={})
+
+    # Normal execute: before -> after returns to baseline.
+    tracker.on_before(conn)
+    assert value() == start + 1
+    assert conn.info[key] is True
+    tracker.on_after(conn)
+    assert value() == start
+    assert key not in conn.info
+
+    # Errored execute: before -> on_error decrements (after never fires).
+    tracker.on_before(conn)
+    assert value() == start + 1
+    tracker.on_error(SimpleNamespace(connection=conn))
+    assert value() == start
+
+    # on_error without a matching before (e.g. connect error) must not push the
+    # gauge negative.
+    tracker.on_error(SimpleNamespace(connection=SimpleNamespace(info={})))
+    assert value() == start