feat(db): add connection retry, adaptive deriver polling, and pool metrics

Add resilience and visibility for DB connection handling under transaction-
pooler (Supavisor) saturation, where client-connection limits get exhausted
across many tenants.

- get_db/tracked_db now force an eager pool checkout with bounded exponential
  backoff (tenacity), retrying SQLAlchemy TimeoutError + OperationalError so
  transient pooler rejections degrade gracefully instead of 500ing. Toggle via
  DB_CONNECTION_RETRY_ENABLED (+ delay/backoff knobs); ~10s default budget.
- Deriver polling backs off when idle or erroring (base -> max, x2 each cycle)
  and snaps back to base on claimed work, cutting steady-state query load.
  Toggle via DERIVER_POLLING_BACKOFF_ENABLED (+ max/multiplier).
- Add scrape-time db_pool_connections Prometheus gauge (checked_out/checked_in/
  size/overflow, labeled api|deriver), registered in both the API lifespan and
  the deriver metrics server.
- Make SqlalchemyIntegration explicit in both Sentry inits; wrap connection
  acquisition in a db.pool.acquire span and capture live pool stats on
  retry-exhaustion.

src/config.py

@@ -622,6 +622,21 @@ class DBSettings(HonchoSettings):
     SQL_DEBUG: bool = False
     TRACING: bool = False
 
+    # Bounded exponential-backoff retry around connection acquisition. Guards
+    # against transient transaction-pooler saturation (e.g. Supavisor rejecting
+    # with "too many clients") by retrying the pool checkout instead of failing
+    # the request immediately. Applied to both the API and background paths.
+    CONNECTION_RETRY_ENABLED: bool = True
+    CONNECTION_RETRY_MAX_DELAY_SECONDS: Annotated[
+        float, Field(default=10.0, gt=0.0, le=120.0)
+    ] = 10.0
+    CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS: Annotated[
+        float, Field(default=0.1, gt=0.0, le=10.0)
+    ] = 0.1
+    CONNECTION_RETRY_BACKOFF_MAX_SECONDS: Annotated[
+        float, Field(default=2.0, gt=0.0, le=30.0)
+    ] = 2.0
+
 
 class AuthSettings(HonchoSettings):
     model_config = SettingsConfigDict(env_prefix="AUTH_", extra="ignore")  # pyright: ignore
@@ -737,6 +752,18 @@ class DeriverSettings(HonchoSettings):
     POLLING_SLEEP_INTERVAL_SECONDS: Annotated[
         float, Field(default=1.0, gt=0.0, le=60.0)
     ] = 1.0
+    # Adaptive polling: when the queue is idle (or the loop is erroring) the
+    # sleep interval grows from POLLING_SLEEP_INTERVAL_SECONDS toward
+    # POLLING_SLEEP_MAX_INTERVAL_SECONDS by POLLING_BACKOFF_MULTIPLIER each
+    # cycle, then snaps back to the base interval as soon as work is found.
+    # Reduces steady-state query load against the (shared) DB/pooler.
+    POLLING_BACKOFF_ENABLED: bool = True
+    POLLING_SLEEP_MAX_INTERVAL_SECONDS: Annotated[
+        float, Field(default=30.0, gt=0.0, le=300.0)
+    ] = 30.0
+    POLLING_BACKOFF_MULTIPLIER: Annotated[
+        float, Field(default=2.0, ge=1.0, le=10.0)
+    ] = 2.0
     STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5
 
     # Retention window (seconds) for keeping errored items in the queue

src/db.py

@@ -1,9 +1,22 @@
 import contextvars
 
+import sentry_sdk
 from sqlalchemy import MetaData, text
-from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.exc import TimeoutError as SQLAlchemyTimeoutError
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
 from sqlalchemy.orm import declarative_base
-from sqlalchemy.pool import NullPool
+from sqlalchemy.pool import NullPool, QueuePool
+from tenacity import (
+    AsyncRetrying,
+    retry_if_exception_type,
+    stop_after_delay,
+    wait_exponential_jitter,
+)
 
 from src.config import settings
 
@@ -45,6 +58,73 @@ SessionLocal = async_sessionmaker(
     bind=engine,
 )
 
+# Errors worth retrying when acquiring a pooled connection: SQLAlchemy's local
+# pool-checkout timeout, and OperationalError (how a saturated transaction
+# pooler surfaces "too many clients" / connection refusals).
+RETRYABLE_DB_CONNECTION_ERRORS = (SQLAlchemyTimeoutError, OperationalError)
+
+
+def get_pool_stats() -> dict[str, int]:
+    """Return live connection-pool stats for this process.
+
+    ``engine.pool`` is the AsyncEngine's pool (the same object as
+    ``engine.sync_engine.pool``); its stat methods are synchronous counter
+    reads with no I/O, so they are safe to call without ``await``. Returns
+    zeros for pools that do not track connections (e.g. ``NullPool``).
+    """
+    zeros = {"checked_out": 0, "checked_in": 0, "size": 0, "overflow": 0}
+    pool = engine.pool
+    # Only QueuePool (and its AsyncAdaptedQueuePool subclass) tracks connection
+    # counts; NullPool and others have no meaningful stats.
+    if not isinstance(pool, QueuePool):
+        return zeros
+    try:
+        return {
+            "checked_out": pool.checkedout(),
+            "checked_in": pool.checkedin(),
+            "size": pool.size(),
+            "overflow": pool.overflow(),
+        }
+    except Exception:
+        return zeros
+
+
+async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None:
+    """Force pool checkout (which ``SessionLocal()`` defers) with bounded backoff.
+
+    ``SessionLocal()`` is lazy: the pool checkout — and any pooler rejection —
+    happens on the first query. We force it here inside a retry block so that
+    transient saturation of the transaction pooler is retried with exponential
+    backoff rather than surfacing as an immediate error. The checkout is wrapped
+    in a Sentry span so wait time is visible in traces; on budget exhaustion the
+    original error is reraised after capturing live pool stats to Sentry.
+
+    Retrying the same session is safe: no connection is bound until checkout
+    succeeds, so each attempt re-attempts the checkout cleanly.
+    """
+    with sentry_sdk.start_span(op="db.pool.acquire", name=context):
+        if not settings.DB.CONNECTION_RETRY_ENABLED:
+            await db.connection()
+            return
+        try:
+            async for attempt in AsyncRetrying(
+                wait=wait_exponential_jitter(
+                    initial=settings.DB.CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS,
+                    max=settings.DB.CONNECTION_RETRY_BACKOFF_MAX_SECONDS,
+                ),
+                stop=stop_after_delay(settings.DB.CONNECTION_RETRY_MAX_DELAY_SECONDS),
+                retry=retry_if_exception_type(RETRYABLE_DB_CONNECTION_ERRORS),
+                reraise=True,
+            ):
+                with attempt:
+                    await db.connection()
+        except RETRYABLE_DB_CONNECTION_ERRORS as e:
+            if settings.SENTRY.ENABLED:
+                sentry_sdk.set_context("db_pool", get_pool_stats())
+                sentry_sdk.capture_exception(e)
+            raise
+
+
 # Define your naming convention
 convention = {
     "ix": "ix_%(table_name)s_%(column_0_N_name)s",  # Index - supports multi-column

src/dependencies.py

@@ -6,7 +6,7 @@ from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src.config import settings
-from src.db import SessionLocal, request_context
+from src.db import SessionLocal, acquire_connection_with_retry, request_context
 
 
 async def get_db():
@@ -16,6 +16,7 @@ async def get_db():
 
     db: AsyncSession = SessionLocal()
     try:
+        await acquire_connection_with_retry(db, context)
         if settings.DB.TRACING:
             await db.execute(
                 text("SELECT set_config('application_name', :name, false)"),
@@ -50,6 +51,7 @@ async def tracked_db(operation_name: str | None = None):
     db = SessionLocal()
 
     try:
+        await acquire_connection_with_retry(db, context or f"task:{operation_name}")
         if settings.DB.TRACING:
             await db.execute(
                 text("SELECT set_config('application_name', :name, false)"),

src/deriver/__main__.py

@@ -8,7 +8,11 @@ from prometheus_client import start_http_server
 from src.config import settings
 from src.db import engine
 from src.startup import validate_embedding_schema
-from src.telemetry import initialize_telemetry_async, shutdown_telemetry
+from src.telemetry import (
+    initialize_telemetry_async,
+    register_db_pool_collector,
+    shutdown_telemetry,
+)
 
 from .queue_manager import main
 
@@ -18,6 +22,8 @@ logger = logging.getLogger(__name__)
 def start_metrics_server() -> None:
     """Start the Prometheus metrics HTTP server on port 9090."""
     start_http_server(9090)
+    # Expose DB connection-pool stats for this deriver instance.
+    register_db_pool_collector("deriver")
     logger.info("Prometheus metrics server started on port 9090")
 
 

src/deriver/queue_manager.py

@@ -11,6 +11,7 @@ import sentry_sdk
 from dotenv import load_dotenv
 from nanoid import generate as generate_nanoid
 from sentry_sdk.integrations.asyncio import AsyncioIntegration
+from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
 from sqlalchemy import and_, delete, or_, select, update
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.engine import CursorResult
@@ -125,6 +126,12 @@ class QueueManager:
         self.worker_ownership: dict[str, WorkerOwnership] = {}
         self.queue_empty_flag: asyncio.Event = asyncio.Event()
 
+        # Current adaptive polling interval; grows while idle/erroring and
+        # resets to the base interval as soon as work is claimed.
+        self._current_poll_interval: float = (
+            settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS
+        )
+
         # Initialize from settings
         self.workers: int = settings.DERIVER.WORKERS
         self.semaphore: asyncio.Semaphore = asyncio.Semaphore(self.workers)
@@ -147,7 +154,9 @@ class QueueManager:
 
         # Initialize Sentry if enabled, using settings
         if settings.SENTRY.ENABLED:
-            initialize_sentry(integrations=[AsyncioIntegration()])
+            initialize_sentry(
+                integrations=[AsyncioIntegration(), SqlalchemyIntegration()]
+            )
 
     def add_task(self, task: asyncio.Task[None]) -> None:
         """Track a new task"""
@@ -378,6 +387,21 @@ class QueueManager:
         )
         return claimed_mapping
 
+    def _reset_poll_interval(self) -> None:
+        """Snap the polling interval back to the base after finding work."""
+        self._current_poll_interval = settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS
+
+    def _advance_poll_interval(self) -> float:
+        """Return the current idle/backoff sleep, then grow it toward the cap."""
+        interval = self._current_poll_interval
+        if settings.DERIVER.POLLING_BACKOFF_ENABLED:
+            self._current_poll_interval = min(
+                self._current_poll_interval
+                * settings.DERIVER.POLLING_BACKOFF_MULTIPLIER,
+                settings.DERIVER.POLLING_SLEEP_MAX_INTERVAL_SECONDS,
+            )
+        return interval
+
     async def polling_loop(self) -> None:
         """Main polling loop to find and process new work units"""
         logger.debug("Starting polling loop")
@@ -385,11 +409,15 @@ class QueueManager:
             while not self.shutdown_event.is_set():
                 if self.queue_empty_flag.is_set():
                     # logger.debug("Queue empty flag set, waiting")
-                    await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS)
+                    # Sleep the already-grown interval; the backoff is advanced
+                    # once per empty-detection below, not here.
+                    await asyncio.sleep(self._current_poll_interval)
                     self.queue_empty_flag.clear()
                     continue
 
-                # Check if we have capacity before querying
+                # Check if we have capacity before querying. There is work to do
+                # (workers are busy), so keep the base interval for fast pickup
+                # when capacity frees rather than backing off.
                 if self.semaphore.locked():
                     # logger.debug("All workers busy, waiting")
                     await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS)
@@ -399,6 +427,7 @@ class QueueManager:
                     await self.cleanup_stale_work_units()
                     claimed_work_units = await self.get_and_claim_work_units()
                     if claimed_work_units:
+                        self._reset_poll_interval()
                         for work_unit_key, aqs_id in claimed_work_units.items():
                             # Create a new task for processing this work unit
                             if not self.shutdown_event.is_set():
@@ -414,15 +443,14 @@ class QueueManager:
                                 self.add_task(task)
                     else:
                         self.queue_empty_flag.set()
-                        await asyncio.sleep(
-                            settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS
-                        )
+                        await asyncio.sleep(self._advance_poll_interval())
                 except Exception as e:
                     logger.exception("Error in polling loop")
                     if settings.SENTRY.ENABLED:
                         sentry_sdk.capture_exception(e)
-                    # Note: rollback is handled by tracked_db dependency
-                    await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS)
+                    # Note: rollback is handled by tracked_db dependency.
+                    # Back off so a down/saturated DB isn't hammered every cycle.
+                    await asyncio.sleep(self._advance_poll_interval())
         finally:
             logger.info("Polling loop stopped")
 

src/main.py

@@ -13,6 +13,7 @@ from fastapi.responses import JSONResponse
 from fastapi_pagination import add_pagination
 from pydantic import ValidationError
 from sentry_sdk.integrations.fastapi import FastApiIntegration
+from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
 from sentry_sdk.integrations.starlette import StarletteIntegration
 
 from src._version import HONCHO_VERSION
@@ -34,6 +35,7 @@ from src.telemetry import (
     initialize_telemetry_async,
     metrics_endpoint,
     prometheus_metrics,
+    register_db_pool_collector,
     shutdown_telemetry,
 )
 from src.telemetry.logging import get_route_template
@@ -117,6 +119,9 @@ if SENTRY_ENABLED:
             FastApiIntegration(
                 transaction_style="endpoint",
             ),
+            # Explicit so DB-query spans (and our db.pool.acquire span) are not
+            # reliant on auto-enabling.
+            SqlalchemyIntegration(),
         ],
         before_send=before_send,
     )
@@ -127,6 +132,9 @@ async def lifespan(_: FastAPI):
     # Initialize CloudEvents telemetry
     await initialize_telemetry_async()
 
+    # Expose DB connection-pool stats for this API instance (no-op if metrics off)
+    register_db_pool_collector("api")
+
     # Validate embedding schema before serving any traffic. Fails closed: if
     # the configured EMBEDDING_VECTOR_DIMENSIONS does not match the physical
     # pgvector columns, the process refuses to start rather than silently

src/telemetry/__init__.py

@@ -12,13 +12,18 @@ This module consolidates all telemetry, metrics, and observability functionality
 """
 
 from src.telemetry.events import emit
-from src.telemetry.prometheus import metrics_endpoint, prometheus_metrics
+from src.telemetry.prometheus import (
+    metrics_endpoint,
+    prometheus_metrics,
+    register_db_pool_collector,
+)
 
 __all__ = [
     "emit",
     "initialize_telemetry_async",
     "metrics_endpoint",
     "prometheus_metrics",
+    "register_db_pool_collector",
     "shutdown_telemetry",
 ]
 

src/telemetry/prometheus/__init__.py

@@ -14,6 +14,7 @@ from src.telemetry.prometheus.metrics import (
     TokenTypes,
     metrics_endpoint,
     prometheus_metrics,
+    register_db_pool_collector,
 )
 
 __all__ = [
@@ -23,4 +24,5 @@ __all__ = [
     "TokenTypes",
     "metrics_endpoint",
     "prometheus_metrics",
+    "register_db_pool_collector",
 ]

src/telemetry/prometheus/metrics.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Iterator
 from enum import Enum
 from typing import cast, final
 
@@ -14,6 +15,7 @@ from prometheus_client import (
     disable_created_metrics,
     generate_latest,
 )
+from prometheus_client.core import GaugeMetricFamily
 from starlette.requests import Request
 from starlette.responses import Response
 
@@ -279,6 +281,46 @@ class PrometheusMetrics:
 prometheus_metrics = PrometheusMetrics()
 
 
+class DBPoolCollector:
+    """Scrape-time collector for SQLAlchemy connection-pool stats.
+
+    Computed live on each /metrics scrape from the async engine's pool, so it
+    is always current with no background task or sampling lag. One instance is
+    registered per process (the API server or a deriver worker).
+    """
+
+    def __init__(self, instance_type: str) -> None:
+        # instance_type: "api" | "deriver"
+        self.instance_type: str = instance_type
+
+    def collect(self) -> Iterator[GaugeMetricFamily]:
+        # Lazy import to avoid an import cycle at module load (db imports config,
+        # telemetry is imported widely). Reads the async engine.pool directly.
+        from src.db import get_pool_stats
+
+        namespace = settings.METRICS.NAMESPACE or ""
+        gauge = GaugeMetricFamily(
+            "db_pool_connections",
+            "DB connections held by this instance, by pool state",
+            labels=["namespace", "instance_type", "state"],
+        )
+        for state, value in get_pool_stats().items():
+            gauge.add_metric([namespace, self.instance_type, state], value)
+        yield gauge
+
+
+_db_pool_collector_registered = False
+
+
+def register_db_pool_collector(instance_type: str) -> None:
+    """Register the DB pool collector once per process (no-op if metrics off)."""
+    global _db_pool_collector_registered
+    if _db_pool_collector_registered or not settings.METRICS.ENABLED:
+        return
+    REGISTRY.register(DBPoolCollector(instance_type))
+    _db_pool_collector_registered = True
+
+
 async def metrics_endpoint(_request: Request) -> Response:
     if not settings.METRICS.ENABLED:
         return Response("Metrics are disabled", status_code=404)

tests/test_dependencies.py

@@ -16,6 +16,12 @@ class FakeSession:
         self.execute_calls: list[tuple[Any, ...]] = []
         self.rollback_calls: int = 0
         self.close_calls: int = 0
+        self.connection_calls: int = 0
+
+    async def connection(self) -> None:
+        # acquire_connection_with_retry forces the (otherwise lazy) pool
+        # checkout via this call before any query runs.
+        self.connection_calls += 1
 
     async def execute(self, statement: Any, params: Any = None) -> None:
         self.execute_calls.append((statement, params))
@@ -44,6 +50,7 @@ async def test_get_db_sets_application_name_when_tracing_enabled(
     try:
         db = await anext(dep_gen)
         assert db is fake_db
+        assert fake_db.connection_calls == 1  # eager pool checkout with retry
         assert len(fake_db.execute_calls) == 1
         stmt, params = fake_db.execute_calls[0]
         assert "set_config" in str(stmt)