From b53b5473585e3217fc361eb76b5c70162ac0236b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 6 Nov 2026 15:50:02 +0800 Subject: [PATCH 1/3] fix(alerts): bound analytics client waits and transaction cleanup --- litellm/repositories/analytics_budget.py | 33 ++++-- .../repositories/test_analytics_budget.py | 101 ++++++++++++++++++ 2 files changed, 128 insertions(+), 6 deletions(-) create mode 100644 tests/unit/repositories/test_analytics_budget.py diff --git a/litellm/repositories/analytics_budget.py b/litellm/repositories/analytics_budget.py index 042bff3fe42..d5e02e655a4 100644 --- a/litellm/repositories/analytics_budget.py +++ b/litellm/repositories/analytics_budget.py @@ -1,6 +1,6 @@ import asyncio -from collections.abc import Generator, Mapping, Sequence -from contextlib import AbstractAsyncContextManager, contextmanager +from collections.abc import AsyncGenerator, Generator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager from contextvars import ContextVar from datetime import timedelta from threading import BoundedSemaphore @@ -75,6 +75,23 @@ def _milliseconds(limit: int) -> int: return limit if remaining is None else min(limit, max(1, int(remaining * 1000))) +@asynccontextmanager +async def _bounded_transaction( + manager: AbstractAsyncContextManager[AnalyticsTransaction], +) -> AsyncGenerator[AnalyticsTransaction, None]: + transaction: Final = await manager.__aenter__() + try: + yield transaction + except BaseException as error: + suppressed: Final = await asyncio.wait_for( + manager.__aexit__(type(error), error, error.__traceback__), timeout=1 + ) + if not suppressed: + raise + else: + await asyncio.wait_for(manager.__aexit__(None, None, None), timeout=1) + + class BoundedAnalyticsDatabase: def __init__(self, database: AnalyticsDatabase, *, background: bool = False) -> None: self._database = database @@ -86,7 +103,9 @@ class BoundedAnalyticsDatabase: raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly") try: remaining: Final = _remaining_seconds() - return await asyncio.wait_for(self._query(query, args), timeout=remaining) + return await asyncio.wait_for( + self._query(query, args), timeout=6.25 if remaining is None else min(6.25, remaining) + ) except RawQueryError as error: if _PostgresError.model_validate(error.meta or {}).code not in ("57014", "55P03"): raise @@ -101,9 +120,11 @@ class BoundedAnalyticsDatabase: self._admission.release() async def _query(self, query: str, args: tuple[object, ...]) -> Sequence[Mapping[str, object]]: - async with self._database.tx( - max_wait=timedelta(milliseconds=_milliseconds(250)), - timeout=timedelta(milliseconds=_milliseconds(6000)), + async with _bounded_transaction( + self._database.tx( + max_wait=timedelta(milliseconds=_milliseconds(250)), + timeout=timedelta(milliseconds=_milliseconds(6000)), + ) ) as transaction: configuration: Final = await transaction.query_raw(_CONFIGURE, self._background, _milliseconds(5000)) if not configuration or not _Admission.model_validate(configuration[0]).acquired: diff --git a/tests/unit/repositories/test_analytics_budget.py b/tests/unit/repositories/test_analytics_budget.py new file mode 100644 index 00000000000..befa645c4b6 --- /dev/null +++ b/tests/unit/repositories/test_analytics_budget.py @@ -0,0 +1,101 @@ +import asyncio +from collections import deque +from collections.abc import AsyncGenerator, Mapping, Sequence +from contextlib import asynccontextmanager, nullcontext +from datetime import timedelta +from typing import Final, Literal + +import pytest + +from litellm.repositories.analytics_budget import ( + AnalyticsCapacityError, + AnalyticsTransaction, + BoundedAnalyticsDatabase, + analytics_deadline, +) + + +class _ClockedLoop(asyncio.SelectorEventLoop): + _ready: deque[asyncio.Handle] + _scheduled: list[asyncio.TimerHandle] + + def __init__(self) -> None: + super().__init__() + self.now = 10000.0 + + def time(self) -> float: + return self.now + + def _run_once(self) -> None: + if not self._ready and self._scheduled: + self.now = max(self.now, self._scheduled[0].when()) + super()._run_once() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # CPython loop hook lacks stubs + + +class _Driver: + def __init__(self, stage: Literal["begin", "query", "commit", "rollback", "healthy"]) -> None: + self.stage = stage + self.cancelled = asyncio.Event() + + async def _stall(self) -> None: + try: + await asyncio.Event().wait() + finally: + self.cancelled.set() + + @asynccontextmanager + async def tx(self, *, max_wait: timedelta, timeout: timedelta) -> AsyncGenerator[AnalyticsTransaction]: + if self.stage == "begin": + await self._stall() + try: + yield self + finally: + if self.stage in ("commit", "rollback"): + await self._stall() + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if query == "SELECT 1 AS result": + if self.stage in ("query", "rollback"): + await self._stall() + return ({"result": 1},) + return ({"acquired": True},) + + +@pytest.mark.parametrize("stage", ("begin", "query", "commit", "rollback")) +@pytest.mark.parametrize("deadline_seconds", (None, 0.25, 10.0)) +def test_stalled_driver_is_cancelled_without_retaining_admission_or_tasks( + stage: Literal["begin", "query", "commit", "rollback"], deadline_seconds: float | None +) -> None: + async def scenario() -> None: + loop: Final = asyncio.get_running_loop() + existing: Final = asyncio.all_tasks() + driver: Final = _Driver(stage) + start: Final = loop.time() + context: Final = nullcontext() if deadline_seconds is None else analytics_deadline(start + deadline_seconds) + with context: + task: Final = asyncio.create_task(BoundedAnalyticsDatabase(driver).query_raw("SELECT 1 AS result")) + done, _ = await asyncio.wait((task,), timeout=10) + if not done: + task.cancel() + watchdog: Final = loop.call_later(1, task.cancel) + try: + with pytest.raises(asyncio.CancelledError): + await task + finally: + watchdog.cancel() + assert done, "Analytics retained its task after the query and cleanup budget" + with pytest.raises(AnalyticsCapacityError): + await task + assert driver.cancelled.is_set() + assert loop.time() - start <= (7.25 if deadline_seconds != 0.25 else 1.25) + healthy: Final = BoundedAnalyticsDatabase(_Driver("healthy")) + assert await asyncio.gather( + healthy.query_raw("SELECT 1 AS result"), healthy.query_raw("SELECT 1 AS result") + ) == [({"result": 1},), ({"result": 1},)] + assert asyncio.all_tasks() == existing + + loop: Final = _ClockedLoop() + try: + loop.run_until_complete(scenario()) + finally: + loop.close() From 8c720dda733a4c0e0d4fe4b5ed9b1a6dabf3af65 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 6 Nov 2026 16:02:05 +0800 Subject: [PATCH 2/3] docs(alerts): explain bounded resource identities --- docs/features/usage-alerts.md | 2 +- tests/integration/database/test_issues.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/features/usage-alerts.md b/docs/features/usage-alerts.md index f9c331b3566..af900e3ffdf 100644 --- a/docs/features/usage-alerts.md +++ b/docs/features/usage-alerts.md @@ -57,7 +57,7 @@ Replace `issue-example` with an ID returned by `/issues`. Activity uses the issu Alert Activity returns exact daily totals in `results` and exact range totals in `metadata`. It does not return the general Usage API's per-day dimension dictionaries. Key issues also return `model_breakdown`, with up to 20 models by range spend by default. Set `model_limit` between 1 and 100 to change that limit. `total_models` counts every matching model; `other` contains the combined metrics and `model_count` of all remaining models, or null when none remain. It is an aggregate, not a model identifier. Adding the displayed models and `other` recovers the full totals. Model identifiers longer than 1,024 characters are included in `other` without truncation. Non-key issues return null for `model_breakdown`. Activity queries use a bounded execution budget. A busy or over-budget query returns `503` with `Retry-After: 5`; retry later or choose a shorter date range. -`/impact` reports distinct resource counts and requests within that same Activity range. These counts mean that resources used the cohort, not that each resource suffered a regression. Saved `impact.measured_keys` and `impact.measured_models`, when present, describe cohorts that met the detector's comparison criteria. Resource pages fetch labels only for the selected page, using the latest deleted-key label when needed. They do not load arbitrary key or user metadata. Display aliases are limited to 256 characters; email labels longer than 320 characters are omitted. Resource IDs remain unchanged. Historical team aggregates do not retain user membership, so a team issue has no `/users` view and its user count is null. User issues similarly have no historical team count. Missing attribution remains unknown. +`/impact` reports distinct resource counts and requests within that same Activity range. These counts mean that resources used the cohort, not that each resource suffered a regression. Saved `impact.measured_keys` and `impact.measured_models`, when present, describe cohorts that met the detector's comparison criteria. Resource pages fetch labels only for the selected page, using the latest deleted-key label when needed. They do not load arbitrary key or user metadata. Display aliases are limited to 256 characters; email labels longer than 320 characters are omitted. Resource pages omit IDs longer than 1,024 characters or 4,096 bytes instead of truncating them. Key IDs that need a `raw:` escape prefix reserve four characters within that limit. Oversized linked user or team IDs are omitted from resource metadata. Activity totals still include traffic from omitted resources. Historical team aggregates do not retain user membership, so a team issue has no `/users` view and its user count is null. User issues similarly have no historical team count. Missing attribution remains unknown. Supporting signals have their own saved comparisons and relationship explanations. Their baseline windows can differ, so do not derive a combined error rate or unit cost by adding their displayed values. Administrators and admin viewers can inspect broader supporting signals. A scoped reader receives only signals for the exact key, user, or team they can access. That restriction applies before the signal count and pagination; broader totals are not disclosed. diff --git a/tests/integration/database/test_issues.py b/tests/integration/database/test_issues.py index 2409ab21a0f..f8c24d12c34 100644 --- a/tests/integration/database/test_issues.py +++ b/tests/integration/database/test_issues.py @@ -945,7 +945,7 @@ async def test_issue_resource_metadata_bounds_labels_and_deleted_history( """INSERT INTO "LiteLLM_DailyUserSpend" (id,user_id,date,api_key,model,custom_llm_provider,api_requests,successful_requests,updated_at) VALUES($1,$2,$3,$4,'example-model',$1,10,10,CURRENT_TIMESTAMP), - ($1||'-wide',repeat(chr(128640),65536),$3,repeat(chr(128640),65536), + ($1||'-wide',repeat(chr(128640),8192),$3,repeat(chr(128640),8192), 'example-model',$1,20,20,CURRENT_TIMESTAMP)""", identity, owner, @@ -976,7 +976,7 @@ async def test_issue_resource_metadata_bounds_labels_and_deleted_history( await _query( db, """UPDATE "LiteLLM_DeletedVerificationToken" - SET user_id=repeat(chr(128640),65536),team_id=repeat(chr(128640),65536) + SET user_id=repeat(chr(128640),8192),team_id=repeat(chr(128640),8192) WHERE id=$1||'-400'""", identity, ) From e48350ae23cb3e0d65194978ecb19b1f93f946c0 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 6 Nov 2026 16:01:41 +0800 Subject: [PATCH 3/3] fix(alerts): retain admission while cancelled database work settles --- .../litellm_core_utils/pending_operations.py | 57 ++++++++ litellm/proxy/db/prisma_client.py | 25 +++- litellm/repositories/analytics_budget.py | 35 ++--- .../test_pending_operations.py | 32 +++++ .../repositories/test_analytics_budget.py | 126 ++++++++++++++++++ 5 files changed, 254 insertions(+), 21 deletions(-) create mode 100644 litellm/litellm_core_utils/pending_operations.py create mode 100644 tests/unit/litellm_core_utils/test_pending_operations.py diff --git a/litellm/litellm_core_utils/pending_operations.py b/litellm/litellm_core_utils/pending_operations.py new file mode 100644 index 00000000000..5c8a242c849 --- /dev/null +++ b/litellm/litellm_core_utils/pending_operations.py @@ -0,0 +1,57 @@ +from collections.abc import Callable, Generator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Final + + +class PendingOperations: + def __init__(self, completed: Callable[[], None]) -> None: + self._completed = completed + self._pending = 0 + self._closed = False + self._released = False + + def started(self) -> bool: + if self._released: + return False + self._pending += 1 + return True + + def finished(self) -> None: + self._pending -= 1 + if self._closed and self._pending == 0: + self._release() + + def close(self) -> None: + self._closed = True + if self._pending == 0: + self._release() + + def _release(self) -> None: + if not self._released: + self._released = True + self._completed() + + +_PENDING: Final[ContextVar[PendingOperations | None]] = ContextVar("litellm_pending_operations", default=None) + + +@contextmanager +def defer_until_operations_finish(completed: Callable[[], None]) -> Generator[None, None, None]: + operations: Final = PendingOperations(completed) + token: Final = _PENDING.set(operations) + try: + yield + finally: + _PENDING.reset(token) + operations.close() + + +def register_pending_operation() -> PendingOperations | None: + operations: Final = _PENDING.get() + return operations if operations is not None and operations.started() else None + + +def finish_pending_operation(operations: PendingOperations | None) -> None: + if operations is not None: + operations.finished() diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index be5de040fdf..6bfb5c7680f 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -19,6 +19,11 @@ from typing import TYPE_CHECKING, Final, Protocol, TypeVar from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.pending_operations import ( + PendingOperations, + finish_pending_operation, + register_pending_operation, +) from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_url_settings import add_missing_query_params, token_refresh_params_from_url from litellm.proxy.db.log_db_metrics import db_io_claimed, record_db_io @@ -148,33 +153,45 @@ class _TrackedPrismaEngine: ) -> _EngineResponseT: self.tracker.begin_operation() task: Final = asyncio.ensure_future(request) + pending: Final = register_pending_operation() try: - return await asyncio.shield(task) + response: Final = await asyncio.shield(task) except asyncio.CancelledError: self._detached = self._detached.union((task,)) - task.add_done_callback(lambda done: self._detached_request_finished(done, abandoned)) + task.add_done_callback(lambda done: self._detached_request_finished(done, abandoned, pending)) raise + except BaseException: + finish_pending_operation(pending) + raise + else: + finish_pending_operation(pending) + return response def _detached_request_finished( self, done: "asyncio.Task[_EngineResponseT]", abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None, + pending: PendingOperations | None, ) -> None: self._detached = self._detached.difference((done,)) if done.cancelled(): + finish_pending_operation(pending) return failure: Final = done.exception() if failure is not None: verbose_proxy_logger.debug("Prisma request abandoned by a cancelled caller failed: %s", failure) + finish_pending_operation(pending) return if abandoned is None: + finish_pending_operation(pending) return settle: Final = asyncio.ensure_future(abandoned(done.result())) self._detached = self._detached.union((settle,)) - settle.add_done_callback(self._abandoned_transaction_settled) + settle.add_done_callback(lambda done: self._abandoned_transaction_settled(done, pending)) - def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]") -> None: + def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]", pending: PendingOperations | None) -> None: self._detached = self._detached.difference((settle,)) + finish_pending_operation(pending) if settle.cancelled() or settle.exception() is None: return verbose_proxy_logger.debug( diff --git a/litellm/repositories/analytics_budget.py b/litellm/repositories/analytics_budget.py index d5e02e655a4..11ac67d970a 100644 --- a/litellm/repositories/analytics_budget.py +++ b/litellm/repositories/analytics_budget.py @@ -9,6 +9,8 @@ from typing import Final, Protocol from prisma.errors import RawQueryError, TransactionExpiredError, TransactionNotStartedError from pydantic import BaseModel +from litellm.litellm_core_utils.pending_operations import defer_until_operations_finish + class AnalyticsTransaction(Protocol): async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... @@ -101,23 +103,22 @@ class BoundedAnalyticsDatabase: async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: if not self._admission.acquire(blocking=False): raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly") - try: - remaining: Final = _remaining_seconds() - return await asyncio.wait_for( - self._query(query, args), timeout=6.25 if remaining is None else min(6.25, remaining) - ) - except RawQueryError as error: - if _PostgresError.model_validate(error.meta or {}).code not in ("57014", "55P03"): - raise - raise AnalyticsCapacityError("Analytics exceeded its query budget. Try again shortly") from error - except ( - TransactionExpiredError, - TransactionNotStartedError, - asyncio.TimeoutError, - ) as error: - raise AnalyticsCapacityError("Analytics exceeded its connection budget. Try again shortly") from error - finally: - self._admission.release() + with defer_until_operations_finish(self._admission.release): + try: + remaining: Final = _remaining_seconds() + return await asyncio.wait_for( + self._query(query, args), timeout=6.25 if remaining is None else min(6.25, remaining) + ) + except RawQueryError as error: + if _PostgresError.model_validate(error.meta or {}).code not in ("57014", "55P03"): + raise + raise AnalyticsCapacityError("Analytics exceeded its query budget. Try again shortly") from error + except ( + TransactionExpiredError, + TransactionNotStartedError, + asyncio.TimeoutError, + ) as error: + raise AnalyticsCapacityError("Analytics exceeded its connection budget. Try again shortly") from error async def _query(self, query: str, args: tuple[object, ...]) -> Sequence[Mapping[str, object]]: async with _bounded_transaction( diff --git a/tests/unit/litellm_core_utils/test_pending_operations.py b/tests/unit/litellm_core_utils/test_pending_operations.py new file mode 100644 index 00000000000..965334bd397 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_pending_operations.py @@ -0,0 +1,32 @@ +import asyncio +from threading import BoundedSemaphore +from typing import Final + +import pytest + +from litellm.litellm_core_utils.pending_operations import ( + defer_until_operations_finish, + finish_pending_operation, + register_pending_operation, +) + + +@pytest.mark.asyncio +async def test_late_child_cannot_release_a_drained_parent_scope_twice() -> None: + admission: Final = BoundedSemaphore(1) + assert admission.acquire(blocking=False) + proceed: Final = asyncio.Event() + + async def late_child() -> None: + await proceed.wait() + operation: Final = register_pending_operation() + assert operation is None + finish_pending_operation(operation) + + with defer_until_operations_finish(admission.release): + child: Final = asyncio.create_task(late_child()) + assert admission.acquire(blocking=False) + proceed.set() + await child + assert not admission.acquire(blocking=False) + admission.release() diff --git a/tests/unit/repositories/test_analytics_budget.py b/tests/unit/repositories/test_analytics_budget.py index befa645c4b6..90f8ab83f3b 100644 --- a/tests/unit/repositories/test_analytics_budget.py +++ b/tests/unit/repositories/test_analytics_budget.py @@ -1,4 +1,5 @@ import asyncio +import heapq from collections import deque from collections.abc import AsyncGenerator, Mapping, Sequence from contextlib import asynccontextmanager, nullcontext @@ -27,6 +28,8 @@ class _ClockedLoop(asyncio.SelectorEventLoop): return self.now def _run_once(self) -> None: + while self._scheduled and self._scheduled[0].cancelled(): + heapq.heappop(self._scheduled) if not self._ready and self._scheduled: self.now = max(self.now, self._scheduled[0].when()) super()._run_once() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # CPython loop hook lacks stubs @@ -99,3 +102,126 @@ def test_stalled_driver_is_cancelled_without_retaining_admission_or_tasks( loop.run_until_complete(scenario()) finally: loop.close() + + +class _Process: + pid: object = None + + +class _HeldEngine: + def __init__(self, stage: Literal["begin", "query", "commit", "rollback"], *, failure: bool = False) -> None: + self.stage = stage + self.failure = failure + self.entered = asyncio.Event() + self.released = asyncio.Event() + self.rollback_entered = asyncio.Event() + self.rollback_released = asyncio.Event() + self.process = _Process() + self.calls = 0 + self.starts = 0 + + async def _request(self, stage: str) -> None: + self.calls += 1 + if stage == self.stage: + self.entered.set() + await self.released.wait() + if self.failure: + raise RuntimeError("synthetic engine error") + + async def query(self, content: str, *, tx_id: str | None) -> object: + if content == "SELECT 1 AS result": + await self._request("query") + return {} + + async def start_transaction(self, *, content: str) -> str: + self.starts += 1 + identifier: Final = f"synthetic-transaction-{self.starts}" + await self._request("begin") + return identifier + + async def commit_transaction(self, tx_id: str) -> None: + await self._request("commit") + + async def rollback_transaction(self, tx_id: str) -> None: + await self._request("rollback") + if self.stage == "begin": + self.rollback_entered.set() + await self.rollback_released.wait() + + +class _TrackedDriver: + def __init__(self, engine: _HeldEngine) -> None: + from litellm.proxy.db.prisma_client import ( + _PrismaDrainTracker, # pyright: ignore[reportPrivateUsage] # exercise real cancellation ownership + _TrackedPrismaEngine, # pyright: ignore[reportPrivateUsage] # exercise real cancellation ownership + ) + + self.stage = engine.stage + self.tracker = _PrismaDrainTracker() + self.engine = _TrackedPrismaEngine(engine, self.tracker) + + @asynccontextmanager + async def tx(self, *, max_wait: timedelta, timeout: timedelta) -> AsyncGenerator[AnalyticsTransaction]: + identifier: Final = await self.engine.start_transaction(content="{}") + try: + yield self + if self.stage == "rollback": + raise RuntimeError("synthetic query failure") + except BaseException: + await self.engine.rollback_transaction(identifier) + raise + else: + await self.engine.commit_transaction(identifier) + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + from litellm.proxy.db.log_db_metrics import claim_db_io + + with claim_db_io(): + await self.engine.query(query, tx_id="synthetic-transaction") + return ({"result": 1},) if query == "SELECT 1 AS result" else ({"acquired": True},) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ("begin", "query", "commit", "rollback")) +@pytest.mark.parametrize("failure", (False, True)) +async def test_cancelled_analytics_keeps_admission_until_detached_engine_work_and_rollback_finish( + stage: Literal["begin", "query", "commit", "rollback"], + failure: bool, +) -> None: + engines: Final = (_HeldEngine(stage, failure=failure), _HeldEngine(stage, failure=failure)) + drivers: Final = tuple(_TrackedDriver(engine) for engine in engines) + tasks: Final = tuple( + asyncio.create_task(BoundedAnalyticsDatabase(driver).query_raw("SELECT 1 AS result")) for driver in drivers + ) + try: + await asyncio.gather(*(engine.entered.wait() for engine in engines)) + for task in tasks: + task.cancel() + for task in tasks: + with pytest.raises(asyncio.CancelledError): + await task + calls: Final = tuple(engine.calls for engine in engines) + for _ in range(64): + with pytest.raises(AnalyticsCapacityError): + await BoundedAnalyticsDatabase(drivers[0]).query_raw("SELECT 1 AS result") + assert tuple(engine.calls for engine in engines) == calls + assert tuple(len(driver.engine._detached) for driver in drivers) == (1, 1) # pyright: ignore[reportPrivateUsage] # bounded retained engine tasks are the behavior under test + for engine in engines: + engine.released.set() + if stage == "begin" and not failure: + await asyncio.gather(*(engine.rollback_entered.wait() for engine in engines)) + with pytest.raises(AnalyticsCapacityError): + await BoundedAnalyticsDatabase(drivers[0]).query_raw("SELECT 1 AS result") + finally: + for engine in engines: + engine.released.set() + engine.rollback_released.set() + await asyncio.gather(*(driver.tracker.wait_until_drained() for driver in drivers)) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert not any(driver.engine._detached for driver in drivers) # pyright: ignore[reportPrivateUsage] # verify detached engine tasks drained + healthy: Final = BoundedAnalyticsDatabase(_Driver("healthy")) + assert await asyncio.gather(healthy.query_raw("SELECT 1 AS result"), healthy.query_raw("SELECT 1 AS result")) == [ + ({"result": 1},), + ({"result": 1},), + ]