diff --git a/config.toml.example b/config.toml.example index 0f4b92c8..258172a7 100644 --- a/config.toml.example +++ b/config.toml.example @@ -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] diff --git a/docs/v2/contributing/configuration.mdx b/docs/v2/contributing/configuration.mdx index c172369c..f2548f94 100644 --- a/docs/v2/contributing/configuration.mdx +++ b/docs/v2/contributing/configuration.mdx @@ -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 diff --git a/docs/v3/contributing/configuration.mdx b/docs/v3/contributing/configuration.mdx index 582bde3a..d7a71b6c 100644 --- a/docs/v3/contributing/configuration.mdx +++ b/docs/v3/contributing/configuration.mdx @@ -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 diff --git a/src/config.py b/src/config.py index c6bfe65e..cf3bb8ce 100644 --- a/src/config.py +++ b/src/config.py @@ -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,10 +623,16 @@ 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. + # Bounded exponential-backoff retry around connection acquisition (used by + # tracked_db for short, DB-only background scopes — NOT the request path). + # 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) @@ -637,6 +644,25 @@ class DBSettings(HonchoSettings): 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 diff --git a/src/db.py b/src/db.py index c543ce26..4fb2c6bd 100644 --- a/src/db.py +++ b/src/db.py @@ -59,12 +59,8 @@ engine = create_async_engine( **engine_kwargs, ) -SessionLocal = async_sessionmaker( - autocommit=False, - autoflush=False, - expire_on_commit=False, - bind=engine, -) +# 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 @@ -99,11 +95,13 @@ def get_pool_stats() -> dict[str, int]: 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": pool.overflow(), + "overflow": max(0, pool.overflow()), } except Exception: return zeros @@ -164,6 +162,93 @@ async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None: _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. 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. + 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 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 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 + + +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. @@ -208,15 +293,20 @@ class DBQueryInflightTracker: _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. Safe to call once per process. + 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 + global _db_instance_type, _inflight_tracker, _db_query_instrumentation_registered _db_instance_type = instance_type - if not settings.METRICS.ENABLED: + 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) @@ -224,6 +314,7 @@ def register_db_query_instrumentation(instance_type: str) -> None: 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 diff --git a/src/dependencies.py b/src/dependencies.py index 710715a7..cd7007e0 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -2,26 +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, acquire_connection_with_retry, request_context +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: - await acquire_connection_with_retry(db, context) - if settings.DB.TRACING: - await db.execute( - text("SELECT set_config('application_name', :name, false)"), - {"name": context}, - ) yield db except Exception: await db.rollback() @@ -31,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 @@ -47,17 +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: - 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)"), - {"name": context or f"task:{operation_name}"}, - ) - yield db except Exception: await db.rollback() diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 2fd141c9..22716cba 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -408,10 +408,9 @@ class QueueManager: try: while not self.shutdown_event.is_set(): if self.queue_empty_flag.is_set(): - # logger.debug("Queue empty flag set, waiting") - # Sleep the already-grown interval; the backoff is advanced - # once per empty-detection below, not here. - await asyncio.sleep(self._current_poll_interval) + # 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 diff --git a/tests/test_db_resilience.py b/tests/test_db_resilience.py index 696a6c53..434e58a6 100644 --- a/tests/test_db_resilience.py +++ b/tests/test_db_resilience.py @@ -9,6 +9,7 @@ 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 @@ -105,6 +106,115 @@ async def test_acquire_disabled_calls_once(monkeypatch: pytest.MonkeyPatch) -> N 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_commit_and_rollback_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) + monkeypatch.setattr(AsyncSession, "commit", noop) + monkeypatch.setattr(AsyncSession, "rollback", noop) + + session = db_module.SessionLocal() + await session.commit() # ensures acquired, commits, then resets + assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage] + + session._honcho_acquired = True # pyright: ignore[reportPrivateUsage] + await session.rollback() + assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage] + + +@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) @@ -133,6 +243,71 @@ def test_polling_backoff_disabled_stays_constant( 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") diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index 511e2a2e..633320b5 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -37,27 +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 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) - 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 @@ -87,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", @@ -97,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 @@ -116,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: @@ -125,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