RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Connection Exponential Backoff (#758) 

* 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.

* 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.

* fix: address CodeRabbit review on PR #758

- db: roll back the session on a retryable checkout failure before
  retrying — a failed autobegin can leave it pending-rollback, making the
  next db.connection() raise instead of re-checking-out cleanly. Cheap
  Python-side cleanup when no connection was bound.
- metrics: guard DBPoolCollector.collect() so a pool-read/import hiccup
  can't raise and abort the whole /metrics scrape (Prometheus drops ALL
  metrics if any collector raises) — log and fall back to empty.

* fix(db): lazy retrying session + review fixes for connection backoff

Address Codex/CodeRabbit review on PR #758.

- Replace eager checkout with HonchoAsyncSession: a lazy AsyncSession that
  checks out its connection (with retry) on the first DB-touching call, not at
  construction. Request handlers doing non-DB work (embedding/file/LLM) before
  their first query no longer pin a connection across it, while the API path
  still gets checkout retry. Only the checkout is retried — the statement runs
  once via super(), so writes are never duplicated. Tracing's set_config moves
  into the same lazy acquire hook.
- Roll the session back on a retryable checkout failure before retrying, so a
  failed autobegin can't leave it pending-rollback.
- Lower default POOL_TIMEOUT to 5s and validate it stays under the retry budget
  for pooled (non-null) POOL_CLASS; update config.toml.example and v2/v3 docs.
- Clamp pool overflow gauge to >= 0 (was negative before the pool fills).
- Remove double-sleep in the deriver idle poll (true backoff cap, not 2x);
  make in-flight instrumentation registration idempotent.
- Tests: HonchoAsyncSession lazy/idempotent acquire, statement-runs-once,
  tracing, commit/rollback flag reset, get_db no-acquire-at-entry, polling-loop
  single-sleep, and the POOL_TIMEOUT/retry-budget validator.

* fix(db): cover all DB-touching session methods; clear flag on close/reset

Address Codex follow-up review on PR #758 (polish, no behavior-critical bug).

- HonchoAsyncSession: wrap get/get_one/stream/stream_scalars/delete in addition
  to execute/scalar/scalars/flush/merge/refresh/commit, so the "lazy checkout
  with retry on first DB use" guarantee has no holes. connection() stays
  unwrapped (acquire_connection_with_retry calls it — wrapping would recurse).
- Reset the acquired flag on close()/reset() too, so a session reused after
  close/reset re-acquires (and re-wraps retry) on its next DB use.
- Fix stale comments: connection retry now applies lazily to the request path
  via HonchoAsyncSession (config.py), and the FakeSession helper note.
- Tests: close/reset flag reset, and get/delete route through acquisition.
This commit is contained in:
Vineeth Voruganti 2026-06-01 12:57:07 -04:00 committed by GitHub
parent 85239a69b2
commit 396976db34
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 924 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
config.toml.example
View File

@ -26,11 +26,18 @@ POOL_CLASS = "default"
POOL_PRE_PING = true
POOL_SIZE = 10
MAX_OVERFLOW = 20
POOL_TIMEOUT = 30 # seconds
POOL_TIMEOUT = 5 # seconds; must stay under CONNECTION_RETRY_MAX_DELAY_SECONDS
POOL_RECYCLE = 300 # seconds
POOL_USE_LIFO = true
SQL_DEBUG = false
TRACING = false
# Bounded retry around connection checkout (used by background/tracked_db and
# the lazy request session). For a pooled (non-null) POOL_CLASS, POOL_TIMEOUT
# must be < CONNECTION_RETRY_MAX_DELAY_SECONDS or startup validation fails.
CONNECTION_RETRY_ENABLED = true
CONNECTION_RETRY_MAX_DELAY_SECONDS = 10.0
CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS = 0.1
CONNECTION_RETRY_BACKOFF_MAX_SECONDS = 2.0
# Authentication settings
[auth]

2
docs/v2/contributing/configuration.mdx
View File

@ -161,7 +161,7 @@ DB_CONNECTION_URI=postgresql+psycopg://honcho_user:secure_password@db.example.co
DB_SCHEMA=public
DB_POOL_SIZE=10
DB_MAX_OVERFLOW=20
DB_POOL_TIMEOUT=30
DB_POOL_TIMEOUT=5
DB_POOL_RECYCLE=300
DB_POOL_PRE_PING=true
DB_SQL_DEBUG=false

8
docs/v3/contributing/configuration.mdx
View File

@ -459,10 +459,16 @@ DB_SCHEMA=public
DB_POOL_PRE_PING=true
DB_POOL_SIZE=10
DB_MAX_OVERFLOW=20
DB_POOL_TIMEOUT=30
DB_POOL_TIMEOUT=5
DB_POOL_RECYCLE=300
DB_POOL_USE_LIFO=true
DB_SQL_DEBUG=false
# Bounded retry around connection checkout. For a pooled (non-null) DB_POOL_CLASS,
# DB_POOL_TIMEOUT must be < DB_CONNECTION_RETRY_MAX_DELAY_SECONDS or startup fails.
DB_CONNECTION_RETRY_ENABLED=true
DB_CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0
DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS=0.1
DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS=2.0
```
### Authentication

59
src/config.py
View File

@ -612,8 +612,9 @@ class DBSettings(HonchoSettings):
POOL_PRE_PING: bool = True
POOL_SIZE: Annotated[int, Field(default=10, gt=0, le=1000)] = 10
MAX_OVERFLOW: Annotated[int, Field(default=20, ge=0, le=1000)] = 20
POOL_TIMEOUT: Annotated[int, Field(default=30, gt=0, le=300)] = (
30 # seconds (max 5 minutes)
POOL_TIMEOUT: Annotated[int, Field(default=5, gt=0, le=300)] = (
5 # seconds; kept under CONNECTION_RETRY_MAX_DELAY_SECONDS so a pool
# checkout fails fast enough to allow a retry within the budget
)
POOL_RECYCLE: Annotated[int, Field(default=300, gt=0, le=7200)] = (
300 # seconds (max 2 hours)
@ -622,6 +623,48 @@ class DBSettings(HonchoSettings):
SQL_DEBUG: bool = False
TRACING: bool = False
# Bounded exponential-backoff retry around connection acquisition. Applied
# lazily on the first DB use of any session (HonchoAsyncSession) — both the
# request path and background/tracked_db scopes — without forcing an eager
# checkout. Guards against transient transaction-pooler saturation (e.g.
# Supavisor rejecting with "too many clients") by retrying the checkout
# instead of failing immediately. CONNECTION_RETRY_MAX_DELAY_SECONDS is the
# TOTAL retry
# budget; with a real QueuePool, a single checkout can block up to
# POOL_TIMEOUT, so POOL_TIMEOUT must stay below the budget for a retry to be
# possible (enforced below). With NullPool (the transaction-pooler setup)
# there is no local queue wait, so saturation is a fast OperationalError and
# POOL_TIMEOUT does not apply.
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
@model_validator(mode="after")
def _validate_retry_budget_vs_pool_timeout(self) -> "DBSettings":
# Only meaningful for a real local pool: NullPool has no queue wait, so
# POOL_TIMEOUT is unused there. For a QueuePool, a checkout can block up
# to POOL_TIMEOUT, so it must be < the total retry budget or the first
# attempt consumes the whole budget and no retry ever happens.
if (
self.POOL_CLASS != "null"
and self.CONNECTION_RETRY_ENABLED
and self.POOL_TIMEOUT >= self.CONNECTION_RETRY_MAX_DELAY_SECONDS
):
raise ValueError(
f"DB_POOL_TIMEOUT ({self.POOL_TIMEOUT}s) must be less than "
+ "DB_CONNECTION_RETRY_MAX_DELAY_SECONDS "
+ f"({self.CONNECTION_RETRY_MAX_DELAY_SECONDS}s) so a pooled checkout "
+ "can fail fast enough to be retried within the budget."
)
return self
class AuthSettings(HonchoSettings):
model_config = SettingsConfigDict(env_prefix="AUTH_", extra="ignore") # pyright: ignore
@ -737,6 +780,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

314
src/db.py
View File

@ -1,11 +1,32 @@
import contextvars
import logging
from typing import Any
from sqlalchemy import MetaData, text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
import sentry_sdk
from sqlalchemy import MetaData, event, text
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
from src.telemetry.prometheus.metrics import (
db_queries_in_flight_gauge,
prometheus_metrics,
)
logger = logging.getLogger(__name__)
connect_args = {"prepare_threshold": None}
@ -38,13 +59,300 @@ engine = create_async_engine(
**engine_kwargs,
)
# NOTE: SessionLocal is defined further down, after HonchoAsyncSession (its
# session class) and acquire_connection_with_retry (which that class calls).
# 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)
# 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.
``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:
# overflow() is negative until the base pool fills (it starts at
# -pool_size); clamp to the count of overflow connections actually open.
return {
"checked_out": pool.checkedout(),
"checked_in": pool.checkedin(),
"size": pool.size(),
"overflow": max(0, 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.
Each attempt rolls the session back on a retryable failure before retrying:
a failed checkout can leave the autobegun transaction in a pending-rollback
state, which would make the next ``db.connection()`` raise instead of
re-checking-out cleanly. The rollback is pure Python-side state cleanup when
no connection was bound, so it is cheap and safe.
"""
with sentry_sdk.start_span(op="db.pool.acquire", name=context):
if not settings.DB.CONNECTION_RETRY_ENABLED:
await db.connection()
return
attempts = 0
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:
attempts += 1
try:
await db.connection()
except RETRYABLE_DB_CONNECTION_ERRORS:
# Reset session state so the next attempt starts clean.
try:
await db.rollback()
except Exception:
logger.debug(
"rollback after failed checkout failed",
exc_info=True,
)
raise
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 HonchoAsyncSession(AsyncSession):
"""AsyncSession that lazily checks out its connection, with retry.
The pool checkout — and any pooler-rejection retry via
``acquire_connection_with_retry`` — happens on the FIRST DB-touching call,
not at construction. So a request handler that does non-DB work (embedding,
file processing, an LLM call) before its first query does NOT pin a
connection across that work, while still getting checkout retry on the
request path.
Only the checkout is retried; the SQL statement itself runs exactly once
(we never retry ``super().execute`` after a broad OperationalError), so
writes are never duplicated. The context for tracing is read from the
``request_context`` ContextVar, which the request/task scope has already set.
"""
# Class-level default; per-instance assignment shadows it (the subclass has
# a __dict__ even though AsyncSession declares __slots__).
_honcho_acquired: bool = False
async def _ensure_acquired(self) -> None:
if self._honcho_acquired:
return
context = request_context.get() or "unknown"
await acquire_connection_with_retry(self, context)
self._honcho_acquired = True
if settings.DB.TRACING:
# Forced checkout already happened above; this rides the same
# connection. super() to avoid re-entering _ensure_acquired.
await super().execute(
text("SELECT set_config('application_name', :name, false)"),
{"name": context},
)
# The overrides below are thin: ensure the connection is checked out (once,
# with retry) before delegating to AsyncSession. They cover every public
# DB-touching async method so the "lazy retry on first DB use" guarantee has
# no holes. Signatures are widened to *args/**kwargs because we only forward;
# call sites are typed against the AsyncSession base, so this does not weaken
# type-checking elsewhere. (connection() is intentionally NOT wrapped —
# acquire_connection_with_retry calls it, so wrapping would recurse.)
async def execute(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().execute(*args, **kwargs)
async def scalar(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().scalar(*args, **kwargs)
async def scalars(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().scalars(*args, **kwargs)
async def get(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().get(*args, **kwargs)
async def get_one(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().get_one(*args, **kwargs)
async def stream(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().stream(*args, **kwargs)
async def stream_scalars(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().stream_scalars(*args, **kwargs)
async def flush(self, *args: Any, **kwargs: Any) -> None:
await self._ensure_acquired()
await super().flush(*args, **kwargs)
async def merge(self, *args: Any, **kwargs: Any) -> Any:
await self._ensure_acquired()
return await super().merge(*args, **kwargs)
async def refresh(self, *args: Any, **kwargs: Any) -> None:
await self._ensure_acquired()
await super().refresh(*args, **kwargs)
async def delete(self, *args: Any, **kwargs: Any) -> None:
await self._ensure_acquired()
await super().delete(*args, **kwargs)
async def commit(self) -> None:
# Ensures the add()->commit() path (autoflush on commit) also retries.
await self._ensure_acquired()
try:
await super().commit()
finally:
# Transaction ended; a later op must re-acquire (and re-wrap retry).
self._honcho_acquired = False
async def rollback(self) -> None:
try:
await super().rollback()
finally:
self._honcho_acquired = False
async def close(self) -> None:
try:
await super().close()
finally:
# The connection is released; a reused session must re-acquire.
self._honcho_acquired = False
async def reset(self) -> None:
try:
await super().reset()
finally:
self._honcho_acquired = False
SessionLocal = async_sessionmaker(
autocommit=False,
autoflush=False,
expire_on_commit=False,
bind=engine,
class_=HonchoAsyncSession,
)
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
_db_query_instrumentation_registered = False
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. Idempotent: repeated calls (e.g. a
re-run lifespan or test startup) won't attach duplicate listeners, which
would double-count in-flight statements.
"""
global _db_instance_type, _inflight_tracker, _db_query_instrumentation_registered
_db_instance_type = instance_type
if not settings.METRICS.ENABLED or _db_query_instrumentation_registered:
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)
_db_query_instrumentation_registered = True
# Define your naming convention
convention = {
"ix": "ix_%(table_name)s_%(column_0_N_name)s", # Index - supports multi-column

34
src/dependencies.py
View File

@ -2,25 +2,22 @@ import uuid
from contextlib import asynccontextmanager
from fastapi import Depends
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession
from src.config import settings
from src.db import SessionLocal, request_context
async def get_db():
"""FastAPI Dependency Generator for Database"""
context = request_context.get() or "unknown"
"""FastAPI Dependency Generator for Database.
The session is lazy: it does NOT check out a pooled connection here.
HonchoAsyncSession acquires it (with retry) on the first DB-touching call,
so a handler doing non-DB work (embedding/file/LLM) before its first query
does not pin a connection across it. Tracing's application_name and the
checkout retry both live in the session, keyed off request_context.
"""
db: AsyncSession = SessionLocal()
try:
if settings.DB.TRACING:
await db.execute(
text("SELECT set_config('application_name', :name, false)"),
{"name": context},
)
yield db
except Exception:
await db.rollback()
@ -30,14 +27,19 @@ async def get_db():
# is closed before the TCP connection drops. Supavisor v2 does NOT
# clean up orphaned transactions on client disconnect in transaction-
# pooling mode, so relying on `in_transaction()` (Python-side state)
# can leave the backend pinned with an open BEGIN.
# can leave the backend pinned with an open BEGIN. (Cheap no-op if the
# lazy session never checked out a connection.)
await db.rollback()
await db.close()
@asynccontextmanager
async def tracked_db(operation_name: str | None = None):
"""Context manager for tracked database sessions"""
"""Context manager for tracked database sessions.
Sets a task-scoped request_context so the lazy session picks it up for
tracing/attribution, then yields a lazy session (see get_db).
"""
# Get request ID if available, or create operation-specific one
context = request_context.get()
token = None
@ -46,16 +48,8 @@ async def tracked_db(operation_name: str | None = None):
context = f"task:{operation_name}:{str(uuid.uuid4())[:8]}"
token = request_context.set(context)
# Create session with tracking info
db = SessionLocal()
try:
if settings.DB.TRACING:
await db.execute(
text("SELECT set_config('application_name', :name, false)"),
{"name": context or f"task:{operation_name}"},
)
yield db
except Exception:
await db.rollback()

11
src/deriver/__main__.py
View File

@ -6,9 +6,13 @@ 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, shutdown_telemetry
from src.telemetry import (
initialize_telemetry_async,
register_db_pool_collector,
shutdown_telemetry,
)
from .queue_manager import main
@ -18,6 +22,9 @@ 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")
register_db_query_instrumentation("deriver")
logger.info("Prometheus metrics server started on port 9090")

45
src/deriver/queue_manager.py
View File

@ -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,18 +387,36 @@ 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")
try:
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)
# The empty-poll branch below already slept this cycle's
# interval; just clear the flag and re-query (no second
# sleep — that would double the effective idle 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 +426,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 +442,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")

11
src/main.py
View File

@ -13,12 +13,13 @@ 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
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,
@ -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,10 @@ 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")
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
# pgvector columns, the process refuses to start rather than silently

7
src/telemetry/__init__.py
View File

@ -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",
]

2
src/telemetry/prometheus/__init__.py
View File

@ -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",
]

79
src/telemetry/prometheus/metrics.py
View File

@ -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
@ -125,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:
@ -275,10 +294,70 @@ 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()
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]:
namespace = settings.METRICS.NAMESPACE or ""
gauge = GaugeMetricFamily(
"db_pool_connections",
"DB connections held by this instance, by pool state",
labels=["namespace", "instance_type", "state"],
)
# Fail soft: Prometheus aborts the entire scrape (dropping ALL metrics)
# if any collector raises, so never let a pool/import hiccup here sink
# the whole /metrics response.
try:
# Lazy import to avoid an import cycle at module load (db imports
# config, telemetry is imported widely). Reads engine.pool directly.
from src.db import get_pool_stats
stats = get_pool_stats()
except Exception:
logger.warning("Failed to collect DB pool stats", exc_info=True)
stats = {}
for state, value in 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)

370
tests/test_db_resilience.py Normal file
View File

@ -0,0 +1,370 @@
"""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
from sqlalchemy.ext.asyncio import AsyncSession
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
self.rollback_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()
async def rollback(self) -> None:
self.rollback_calls += 1
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
# Session is reset after each failed checkout before the next attempt.
assert session.rollback_calls == 2
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
# --- HonchoAsyncSession: lazy checkout with retry on first DB use ------------
@pytest.mark.asyncio
async def test_session_lazy_acquires_on_first_db_use(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""No checkout at construction; acquisition (with retry) happens once on the
first execute, and the statement itself runs exactly once."""
monkeypatch.setattr(settings.DB, "TRACING", False)
acquired: list[str] = []
executed: list[Any] = []
async def fake_acquire(_session: Any, context: str) -> None:
acquired.append(context)
async def fake_super_execute(_self: Any, *args: Any, **_kw: Any) -> str:
executed.append(args[0] if args else None)
return "result"
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
monkeypatch.setattr(AsyncSession, "execute", fake_super_execute)
session = db_module.SessionLocal()
assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage]
assert acquired == [] # constructing the session does NOT check out
result = await session.execute("SELECT 1")
assert result == "result"
assert acquired == ["unknown"] # acquired exactly once, on first use
assert session._honcho_acquired is True # pyright: ignore[reportPrivateUsage]
assert executed == ["SELECT 1"] # statement ran once (not retried)
await session.execute("SELECT 2")
assert acquired == ["unknown"] # still once — no re-acquire on later use
@pytest.mark.asyncio
async def test_session_tracing_sets_application_name_on_acquire(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(settings.DB, "TRACING", True)
calls: list[tuple[Any, ...]] = []
async def fake_acquire(_session: Any, _context: str) -> None:
return None
async def fake_super_execute(_self: Any, *args: Any, **_kw: Any) -> None:
calls.append(args)
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
monkeypatch.setattr(AsyncSession, "execute", fake_super_execute)
token = db_module.request_context.set("request:trace-ctx")
try:
session = db_module.SessionLocal()
await session.execute("SELECT 1")
finally:
db_module.request_context.reset(token)
# First super().execute is the set_config, then the real statement.
assert any("set_config" in str(c[0]) for c in calls)
set_config_call = next(c for c in calls if "set_config" in str(c[0]))
assert set_config_call[1] == {"name": "request:trace-ctx"}
@pytest.mark.asyncio
async def test_session_lifecycle_methods_reset_acquired_flag(
monkeypatch: pytest.MonkeyPatch,
) -> None:
async def fake_acquire(_session: Any, _context: str) -> None:
return None
async def noop(_self: Any) -> None:
return None
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
for method in ("commit", "rollback", "close", "reset"):
monkeypatch.setattr(AsyncSession, method, noop)
session = db_module.SessionLocal()
await session.commit() # ensures acquired, commits, then resets
assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage]
# rollback/close/reset must each clear the flag so a reused session
# re-acquires (and re-wraps retry) on its next DB use.
for method in ("rollback", "close", "reset"):
session._honcho_acquired = True # pyright: ignore[reportPrivateUsage]
await getattr(session, method)()
assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage]
@pytest.mark.asyncio
async def test_session_get_and_delete_also_acquire(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The lazy-acquire guarantee covers get/delete, not just execute."""
monkeypatch.setattr(settings.DB, "TRACING", False)
acquired: list[str] = []
async def fake_acquire(_session: Any, context: str) -> None:
acquired.append(context)
async def fake_get(_self: Any, *_a: Any, **_k: Any) -> str:
return "row"
async def fake_delete(_self: Any, *_a: Any, **_k: Any) -> None:
return None
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
monkeypatch.setattr(AsyncSession, "get", fake_get)
monkeypatch.setattr(AsyncSession, "delete", fake_delete)
session = db_module.SessionLocal()
await session.get(object, 1)
await session.delete(object())
assert acquired == ["unknown"] # acquired once on the first DB-touching call
@pytest.mark.asyncio
async def test_get_db_does_not_acquire_at_dependency_entry(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.dependencies import get_db as real_get_db
acquired: list[str] = []
async def fake_acquire(_session: Any, context: str) -> None:
acquired.append(context)
monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire)
dep_gen = real_get_db()
await anext(dep_gen)
assert acquired == [] # yielding the session must not check out
await dep_gen.aclose()
assert acquired == [] # finally rollback/close must not check out either
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_pool_timeout_must_be_under_retry_budget() -> None:
"""A pooled QueuePool checkout that can block past the retry budget is a
contradiction (no retry ever happens) and must fail config validation."""
from src.config import DBSettings
# Contradiction: pooled + retry on, but POOL_TIMEOUT >= budget.
with pytest.raises(ValueError, match="DB_POOL_TIMEOUT"):
DBSettings(
POOL_CLASS="default",
POOL_TIMEOUT=30,
CONNECTION_RETRY_ENABLED=True,
CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0,
)
# NullPool has no queue wait, so POOL_TIMEOUT is irrelevant -> allowed.
DBSettings(
POOL_CLASS="null",
POOL_TIMEOUT=30,
CONNECTION_RETRY_ENABLED=True,
CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0,
)
@pytest.mark.asyncio
async def test_polling_loop_idle_sleeps_once_per_cycle(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Drive the real loop on an empty queue: exactly one (growing, capped)
sleep per empty poll — no double-sleep from the queue_empty_flag branch."""
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", 8.0)
import asyncio
from src.deriver import queue_manager as qm_mod
qm = qm_mod.QueueManager()
sleeps: list[float] = []
polls = {"n": 0}
async def fake_cleanup() -> None:
return None
async def fake_claim() -> dict[str, str]:
polls["n"] += 1
if polls["n"] >= 5:
qm.shutdown_event.set() # stop after 5 empty polls
return {}
async def fake_sleep(seconds: float) -> None:
sleeps.append(seconds)
monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup)
monkeypatch.setattr(qm, "get_and_claim_work_units", fake_claim)
# queue_manager calls asyncio.sleep on the stdlib module; patch it there.
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
await qm.polling_loop()
# One sleep per empty poll, growing 1->2->4->8 then capped at 8 (not doubled).
assert sleeps == [1.0, 2.0, 4.0, 8.0, 8.0]
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

33
tests/test_dependencies.py
View File

@ -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:
# Tracks checkout attempts so tests can assert get_db/tracked_db stay
# lazy (they should never force a checkout themselves).
self.connection_calls += 1
async def execute(self, statement: Any, params: Any = None) -> None:
self.execute_calls.append((statement, params))
@ -31,26 +37,24 @@ class FakeSession:
@pytest.mark.asyncio
async def test_get_db_sets_application_name_when_tracing_enabled(
async def test_get_db_yields_lazily_without_checkout_or_tracing(
monkeypatch: pytest.MonkeyPatch,
) -> None:
# get_db must NOT touch the connection or run set_config itself — those now
# happen lazily inside HonchoAsyncSession on first DB use, so a handler doing
# non-DB work before its first query never pins a connection.
fake_db = FakeSession()
monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
monkeypatch.setattr(settings.DB, "TRACING", True)
monkeypatch.setattr(settings.DB, "TRACING", True) # still no set_config here
context_token = request_context.set("request:test-ctx")
dep_gen = real_get_db()
try:
db = await anext(dep_gen)
assert db is fake_db
assert len(fake_db.execute_calls) == 1
stmt, params = fake_db.execute_calls[0]
assert "set_config" in str(stmt)
assert params == {"name": "request:test-ctx"}
assert fake_db.connection_calls == 0 # no eager checkout
assert fake_db.execute_calls == [] # no set_config in get_db
finally:
await dep_gen.aclose()
request_context.reset(context_token)
assert fake_db.rollback_calls == 1 # unconditional rollback in finally
assert fake_db.close_calls == 1
@ -80,7 +84,6 @@ async def test_tracked_db_creates_and_resets_task_context(
) -> None:
fake_db = FakeSession()
monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
monkeypatch.setattr(settings.DB, "TRACING", True)
monkeypatch.setattr(
uuid,
"uuid4",
@ -90,15 +93,12 @@ async def test_tracked_db_creates_and_resets_task_context(
clear_token = request_context.set(None)
try:
async with real_tracked_db("cleanup_job"):
# tracked_db sets the task context so the lazy session can read it.
assert request_context.get() == "task:cleanup_job:12345678"
finally:
request_context.reset(clear_token)
assert request_context.get() is None
assert len(fake_db.execute_calls) == 1
stmt, params = fake_db.execute_calls[0]
assert "set_config" in str(stmt)
assert params == {"name": "task:cleanup_job:12345678"}
assert fake_db.rollback_calls == 1 # unconditional rollback in finally
assert fake_db.close_calls == 1
@ -109,7 +109,6 @@ async def test_tracked_db_preserves_existing_request_context(
) -> None:
fake_db = FakeSession()
monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
monkeypatch.setattr(settings.DB, "TRACING", True)
context_token = request_context.set("request:existing")
try:
@ -118,10 +117,6 @@ async def test_tracked_db_preserves_existing_request_context(
finally:
request_context.reset(context_token)
assert len(fake_db.execute_calls) == 1
stmt, params = fake_db.execute_calls[0]
assert "set_config" in str(stmt)
assert params == {"name": "request:existing"}
assert fake_db.rollback_calls == 1 # unconditional rollback in finally
assert fake_db.close_calls == 1