diff --git a/src/config.py b/src/config.py index cf3bb8ce..35d72d3e 100644 --- a/src/config.py +++ b/src/config.py @@ -623,11 +623,13 @@ class DBSettings(HonchoSettings): SQL_DEBUG: bool = False TRACING: bool = False - # 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 + # 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) diff --git a/src/db.py b/src/db.py index 4fb2c6bd..0a133c05 100644 --- a/src/db.py +++ b/src/db.py @@ -197,9 +197,12 @@ class HonchoAsyncSession(AsyncSession): ) # 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. + # 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) @@ -212,6 +215,22 @@ class HonchoAsyncSession(AsyncSession): 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) @@ -224,6 +243,10 @@ class HonchoAsyncSession(AsyncSession): 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() @@ -239,6 +262,19 @@ class HonchoAsyncSession(AsyncSession): 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, diff --git a/tests/test_db_resilience.py b/tests/test_db_resilience.py index 434e58a6..a82f307a 100644 --- a/tests/test_db_resilience.py +++ b/tests/test_db_resilience.py @@ -173,7 +173,7 @@ async def test_session_tracing_sets_application_name_on_acquire( @pytest.mark.asyncio -async def test_session_commit_and_rollback_reset_acquired_flag( +async def test_session_lifecycle_methods_reset_acquired_flag( monkeypatch: pytest.MonkeyPatch, ) -> None: async def fake_acquire(_session: Any, _context: str) -> None: @@ -183,16 +183,46 @@ async def test_session_commit_and_rollback_reset_acquired_flag( return None monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire) - monkeypatch.setattr(AsyncSession, "commit", noop) - monkeypatch.setattr(AsyncSession, "rollback", noop) + 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] - session._honcho_acquired = True # pyright: ignore[reportPrivateUsage] - await session.rollback() - 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 diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index 633320b5..2996cddc 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -19,8 +19,8 @@ class FakeSession: 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. + # 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: