diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 7f8db72a1a6..ef4559174ce 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -44,12 +44,13 @@ import sys import tempfile import urllib.parse from collections.abc import Callable, Mapping, Sequence +from datetime import timedelta from functools import partial from pathlib import Path from types import MappingProxyType from typing import Annotated, Final, Protocol, TypeAlias, cast -from pydantic import AliasChoices, BeforeValidator, Field +from pydantic import AliasChoices, BeforeValidator, Field, TypeAdapter, ValidationError from pydantic_settings import BaseSettings, SettingsConfigDict from litellm.proxy.db.pgbouncer import database_url_is_pooled @@ -300,6 +301,23 @@ def connection_params_from_url(url: str) -> Mapping[str, str | int | float]: ) +PRISMA_DEFAULT_POOL_TIMEOUT: Final = timedelta(seconds=10) +_POOL_TIMEOUT_SECONDS: Final = TypeAdapter(float) + + +def pool_timeout_from_url(url: str) -> timedelta: + """How long Prisma waits for a pooled connection on ``url``: its ``pool_timeout`` param, else Prisma's default. + + ``pool_timeout=0`` tells Prisma to wait forever, which a transaction's ``max_wait`` cannot express, so it falls + back to the default too. + """ + try: + seconds: Final = _POOL_TIMEOUT_SECONDS.validate_python(connection_params_from_url(url).get("pool_timeout")) + except ValidationError: + return PRISMA_DEFAULT_POOL_TIMEOUT + return timedelta(seconds=seconds) if seconds > 0 else PRISMA_DEFAULT_POOL_TIMEOUT + + # A re-minted token URL replaces a URL to the same database, so unlike the # reader allowlist it may also carry ``options``: that is where the server-side # timeouts (statement, lock, idle-in-transaction) live, and a refresh that diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 8fbc8708f8c..3cddd84e85c 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -30,11 +30,12 @@ _CONNECTION_CAPACITY_PHRASES: Final = ( "too many connections for role", "too many connections for database", "remaining connection slots are reserved", + "timed out fetching a new connection from the connection pool", + "unable to start a transaction in the given time", ) _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" -_STATEMENT_TIMEOUT_SQLSTATE: Final = "57014" _STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout" @@ -274,10 +275,11 @@ class PrismaDBExceptionHandler: def is_database_capacity_error(e: Exception) -> bool: """True iff Postgres refused the pool a new connection (SQLSTATE 53300: ``too many clients already``, a reserved slot, or a per-role limit) or - prisma's own pool timed out handing one over (P2024). The server is up - but full, so the failure is neither a transport error (which would tear - the engine down and open yet more connections against it) nor a - rejection of the rows being written.""" + prisma's own pool timed out handing one over (P2024 for a query, P2028 + ``Unable to start a transaction in the given time`` for a transaction). + The server is up but full, so the failure is neither a transport error + (which would tear the engine down and open yet more connections against + it) nor a rejection of the rows being written.""" import prisma if not isinstance(e, _exception_types(prisma.errors.PrismaError)): @@ -306,13 +308,14 @@ class PrismaDBExceptionHandler: @staticmethod def is_statement_timeout_error(e: Exception) -> bool: - """True iff Postgres cancelled a statement for exceeding ``statement_timeout`` (SQLSTATE 57014).""" + """True iff Postgres cancelled a statement for exceeding ``statement_timeout``. + + Only the message decides: its SQLSTATE 57014 is the same one + ``pg_cancel_backend`` and a client-side cancel raise.""" import prisma if not isinstance(e, _exception_types(prisma.errors.PrismaError)): return False - if PrismaDBExceptionHandler.postgres_sqlstate(e) == _STATEMENT_TIMEOUT_SQLSTATE: - return True return _STATEMENT_TIMEOUT_PHRASE in str(e).lower() @staticmethod diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index be5de040fdf..a660ee0df97 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -20,7 +20,11 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger 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.db_url_settings import ( + add_missing_query_params, + pool_timeout_from_url, + token_refresh_params_from_url, +) from litellm.proxy.db.log_db_metrics import db_io_claimed, record_db_io from litellm.proxy.db.prisma_query_span import parse_prisma_query from litellm.proxy.db.token_auth import ( @@ -904,6 +908,11 @@ class PrismaWrapper: async def disconnect(self, timeout: float | timedelta | None = None) -> None: ... + @property + def pool_timeout(self) -> timedelta: + """How long this wrapper's engine waits for a pooled connection, read from its URL's ``pool_timeout``.""" + return pool_timeout_from_url(os.getenv(self._db_url_env_var, "")) + def __getattr__(self, name: str): """ Proxy attribute access to the underlying Prisma client. diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e68d348c16c..8e4724629c9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5315,6 +5315,7 @@ _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) _PRISMA_DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2) _SPEND_LOG_READ_STATEMENT_TIMEOUT: Final = timedelta(seconds=SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS) _STATEMENT_TIMEOUT_TRANSACTION_GRACE: Final = timedelta(seconds=5) +_STATEMENT_TIMEOUT_MS: Final = TypeAdapter(int) class _ReadTransactionClient(Protocol): @@ -5322,6 +5323,27 @@ class _ReadTransactionClient(Protocol): def read_db(self) -> PrismaWrapper: ... +class _RawReadTransaction(Protocol): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + async def execute_raw(self, query: str, *args: object) -> int: ... + + +async def _lower_statement_timeout(transaction: _RawReadTransaction, timeout_ms: int) -> int: + """``SET LOCAL statement_timeout`` to ``timeout_ms`` unless the session already runs under a stricter one. + + Returns the limit the read runs under, in milliseconds. 0 means the session has none. + """ + rows: Final = await transaction.query_raw( + "SELECT setting::int AS ms FROM pg_settings WHERE name = 'statement_timeout'" + ) + current_ms: Final = _STATEMENT_TIMEOUT_MS.validate_python(next(iter(rows), {}).get("ms", 0)) + if 0 < current_ms <= timeout_ms: + return current_ms + await transaction.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") + return timeout_ms + + async def query_raw_within_statement_timeout( prisma_client: _ReadTransactionClient, sql: str, @@ -5333,19 +5355,24 @@ async def query_raw_within_statement_timeout( Prisma's HTTP client gives up on a statement after 30s anyway, so without this the Admin UI's error is the same but Postgres keeps running the scan. - A cancelled read surfaces as a 503 asking for a narrower date range. + A cancelled read surfaces as a 503 asking for a narrower date range. The + transaction waits for a pooled connection as long as a plain read would, + and never loosens a ``statement_timeout`` the role or database already set. """ - try: - async with prisma_client.read_db.tx(timeout=timeout + _STATEMENT_TIMEOUT_TRANSACTION_GRACE) as transaction: - await transaction.execute_raw(f"SET LOCAL statement_timeout = {int(timeout / timedelta(milliseconds=1))}") + read_db: Final = prisma_client.read_db + async with read_db.tx( + timeout=timeout + _STATEMENT_TIMEOUT_TRANSACTION_GRACE, max_wait=read_db.pool_timeout + ) as transaction: + effective_ms: Final = await _lower_statement_timeout(transaction, int(timeout / timedelta(milliseconds=1))) + try: return await transaction.query_raw(sql, *params) - except Exception as e: - if not PrismaDBExceptionHandler.is_statement_timeout_error(e): - raise - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail={"error": f"{what} took longer than {int(timeout.total_seconds())}s. Pick a shorter date range"}, - ) from e + except Exception as e: + if not PrismaDBExceptionHandler.is_statement_timeout_error(e): + raise + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"error": f"{what} took longer than {effective_ms / 1000:g}s. Pick a shorter date range"}, + ) from e async def _lookup_deprecated_key( diff --git a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py index 2d28292a487..05606f17e1d 100644 --- a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py +++ b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py @@ -77,7 +77,7 @@ def build_prisma(query_raw: AsyncMock) -> MagicMock: def _spend_log_reads(prisma: MagicMock) -> AsyncMock: """The raw reads the cache dashboard ran under the statement timeout.""" - return prisma.read_db.tx.return_value.query_raw + return prisma.read_db.tx.return_value.reads def dispatching_query_raw() -> AsyncMock: diff --git a/tests/unit/proxy/db/fake_prisma_engine.py b/tests/unit/proxy/db/fake_prisma_engine.py index 082659a3adc..8947a8a349d 100644 --- a/tests/unit/proxy/db/fake_prisma_engine.py +++ b/tests/unit/proxy/db/fake_prisma_engine.py @@ -2,6 +2,7 @@ marking the DB I/O witness the way ``_TrackedPrismaEngine`` does, so the producer under test emits its service event.""" +from datetime import timedelta from typing import Final, TypeVar from unittest.mock import AsyncMock, MagicMock @@ -21,11 +22,20 @@ def engine_call(return_value: _T | None = None) -> AsyncMock: class RecordingReadTransaction: - """The transaction ``query_raw_within_statement_timeout`` opens on ``read_db``, keeping its ``SET LOCAL`` statements.""" + """The transaction ``query_raw_within_statement_timeout`` opens on ``read_db``, keeping its ``SET LOCAL`` statements. - def __init__(self, query_raw: AsyncMock) -> None: + The helper's ``pg_settings`` probe sees ``statement_timeout_ms`` as the session's current ``statement_timeout`` + (0 means unlimited), every other raw read is answered by ``reads``.""" + + def __init__(self, reads: AsyncMock, statement_timeout_ms: int = 0) -> None: self.statements: tuple[str, ...] = () - self.query_raw: AsyncMock = query_raw + self.reads: AsyncMock = reads + self.statement_timeout_ms: Final = statement_timeout_ms + + async def query_raw(self, sql: str, *args: object) -> object: + if "FROM pg_settings" in sql: + return [{"ms": self.statement_timeout_ms}] + return await self.reads(sql, *args) async def execute_raw(self, sql: str) -> int: self.statements = (*self.statements, sql) @@ -38,9 +48,11 @@ class RecordingReadTransaction: pass -def statement_timeout_reads(prisma: MagicMock, query_raw: AsyncMock) -> RecordingReadTransaction: +def statement_timeout_reads( + prisma: MagicMock, query_raw: AsyncMock, statement_timeout_ms: int = 0 +) -> RecordingReadTransaction: """Make ``prisma.read_db.tx()`` hand out one recording transaction whose raw reads answer with ``query_raw``.""" - transaction: Final = RecordingReadTransaction(query_raw) + transaction: Final = RecordingReadTransaction(query_raw, statement_timeout_ms) prisma.read_db.tx = MagicMock(return_value=transaction) return transaction @@ -71,6 +83,8 @@ class _DelegatingReadTransaction: return 0 async def query_raw(self, sql: str, *args: object) -> object: + if "FROM pg_settings" in sql: + return [{"ms": 0}] return await getattr(self.db, "query_raw")(sql, *args) async def __aenter__(self) -> "_DelegatingReadTransaction": @@ -83,6 +97,8 @@ class _DelegatingReadTransaction: class ReadTargetOf: """What a prisma fake returns from ``read_db``: ``tx()`` opens a transaction over the fake's ``db``.""" + pool_timeout: Final = timedelta(seconds=10) + def __init__(self, db: object) -> None: self.db: Final = db diff --git a/tests/unit/proxy/db/test_db_url_settings.py b/tests/unit/proxy/db/test_db_url_settings.py index 99957447248..dd820377458 100644 --- a/tests/unit/proxy/db/test_db_url_settings.py +++ b/tests/unit/proxy/db/test_db_url_settings.py @@ -33,7 +33,9 @@ from pydantic import ValidationError from litellm.proxy.db.db_url_settings import ( PG_SSL_REQUEST, + PRISMA_DEFAULT_POOL_TIMEOUT, DatabaseURLSettings, + pool_timeout_from_url, token_refresh_params_from_url, translate_libpq_ssl_params, unsupported_db_scheme, @@ -1314,3 +1316,18 @@ def test_connection_budget_message_survives_an_unparseable_limit() -> None: message: Final = postgres_connection_budget_message(writer_limit="ten", reader_limit=None, num_workers="2") assert "writer='ten'" in message + + +@pytest.mark.parametrize( + ("url", "expected"), + [ + ("postgresql://u:p@db:5432/litellm?connection_limit=10&pool_timeout=60", datetime.timedelta(seconds=60)), + ("postgresql://u:p@db:5432/litellm?pool_timeout=2.5", datetime.timedelta(seconds=2.5)), + ("postgresql://u:p@db:5432/litellm?connection_limit=10", PRISMA_DEFAULT_POOL_TIMEOUT), + ("postgresql://u:p@db:5432/litellm?pool_timeout=0", PRISMA_DEFAULT_POOL_TIMEOUT), + ("postgresql://u:p@db:5432/litellm?pool_timeout=soon", PRISMA_DEFAULT_POOL_TIMEOUT), + ], + ids=["set", "fractional", "absent", "zero_waits_forever", "garbage"], +) +def test_pool_timeout_from_url(url: str, expected: datetime.timedelta) -> None: + assert pool_timeout_from_url(url) == expected diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index 46704bf52ef..bb27c73b50c 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -842,11 +842,27 @@ def _pool_timeout_error() -> DataError: ) +def _transaction_start_timeout_error() -> EngineRequestError: + """The shape prisma raises when ``tx()`` could not get a pooled connection within ``max_wait`` (P2028).""" + return EngineRequestError( + MagicMock(status=504), + '{"is_panic":false,"message":"Transaction API error: Unable to start a transaction in the given time.",' + '"error_code":"P2028"}', + ) + + @pytest.mark.parametrize( "error", [ _capacity_error("sorry, too many clients already"), _pool_timeout_error(), + _transaction_start_timeout_error(), + EngineRequestError( + MagicMock(status=500), + '{"is_panic":false,"message":"Timed out fetching a new connection from the connection pool. More info: ' + 'http://pris.ly/d/connection-pool (Current connection pool timeout: 1, connection limit: 10)",' + '"error_code":"P2024"}', + ), EngineRequestError( MagicMock(status=500), '{"is_panic":false,"message":"Error in connector: Error querying the database: FATAL: sorry, too many clients already","backtrace":null}', @@ -861,7 +877,14 @@ def _pool_timeout_error() -> DataError: } ), ], - ids=["53300_connector_dataerror", "P2024_pool_timeout", "engine_500", "raw_query_53300"], + ids=[ + "53300_connector_dataerror", + "P2024_pool_timeout", + "P2028_transaction_start", + "P2024_inside_transaction", + "engine_500", + "raw_query_53300", + ], ) def test_is_database_capacity_error_recognises_postgres_and_pool_exhaustion(error: Exception) -> None: """SQLSTATE 53300 and prisma P2024 mean the statement was never sent because @@ -958,11 +981,21 @@ def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None [ (_lock_timeout_error(message="canceling statement due to statement timeout", code="57014"), True), (_lock_timeout_error(message="canceling statement due to statement timeout", code=None), True), + (_lock_timeout_error(message="canceling statement due to user request", code="57014"), False), (_lock_timeout_error(), False), (_capacity_error("sorry, too many clients already"), False), + (_transaction_start_timeout_error(), False), (RuntimeError("canceling statement due to statement timeout"), False), ], - ids=["meta_57014", "message_only", "lock_timeout_55P03", "capacity_53300", "not_prisma"], + ids=[ + "meta_57014", + "message_only", + "pg_cancel_backend_57014", + "lock_timeout_55P03", + "capacity_53300", + "P2028_transaction_start", + "not_prisma", + ], ) def test_is_statement_timeout_error(error: Exception, expected: bool) -> None: assert PrismaDBExceptionHandler.is_statement_timeout_error(error) is expected diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 6c82628576f..223042aab76 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/tests/unit/proxy/db/test_prisma_client.py @@ -602,3 +602,16 @@ def test_migration_dry_run_summary_needs_database_url(unset_database_url): from litellm.proxy.db.prisma_client import PrismaManager assert PrismaManager.migration_dry_run_summary(use_migrate=True) == (False, "DATABASE_URL is not set.") + + +def test_prisma_wrapper_pool_timeout_comes_from_its_own_database_url(monkeypatch: pytest.MonkeyPatch) -> None: + from datetime import timedelta + + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/litellm?pool_timeout=60") + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@replica:5432/litellm?pool_timeout=7") + + wrapper: Final = PrismaWrapper( + original_prisma=MagicMock(), iam_token_db_auth=False, db_url_env_var="DATABASE_URL_READ_REPLICA" + ) + + assert wrapper.pool_timeout == timedelta(seconds=7) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 697d9810284..5375c0878a6 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -17307,7 +17307,7 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m assert any("statement_timeout = 30000" in statement for statement in transaction.statements) mock_db_client.db.tx.assert_not_called() - sql, *params = transaction.query_raw.call_args.args + sql, *params = transaction.reads.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"] assert 'FROM "LiteLLM_SpendLogs" sl' in sql assert 'sl."startTime" >= $1::timestamp' in sql @@ -17343,7 +17343,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" ) - sql, *params = transaction.query_raw.call_args.args + sql, *params = transaction.reads.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha"] assert 'sl."user" = $' not in sql @@ -17367,7 +17367,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" ) - sql, *params = transaction.query_raw.call_args.args + sql, *params = transaction.reads.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"] assert "sl.team_id IN ($3)" in sql assert 'AND sl."user" = $4' in sql @@ -17390,7 +17390,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie ) assert exc_info.value.status_code == 404 - transaction.query_raw.assert_not_called() + transaction.reads.assert_not_called() @pytest.mark.asyncio @@ -17418,7 +17418,7 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids assert exc_info.value.status_code == 400 assert expected_error in str(exc_info.value.detail) - transaction.query_raw.assert_not_called() + transaction.reads.assert_not_called() @pytest.mark.asyncio diff --git a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py index aaf4bc5438c..e03eef260a5 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -205,7 +205,7 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): def _spend_log_reads(mock_prisma: MagicMock) -> AsyncMock: """The raw reads ``/spend/logs/ui`` ran under the statement timeout.""" - return mock_prisma.read_db.tx.return_value.query_raw + return mock_prisma.read_db.tx.return_value.reads def _make_ui_spend_logs_mock(count_total, page_rows): diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 698f4787a42..48ff1417c7a 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -3734,10 +3734,67 @@ async def test_query_raw_within_statement_timeout_has_postgres_cancel_the_read_a assert rows == [{"spend": 1.5}] assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) - transaction.query_raw.assert_awaited_once_with('SELECT spend FROM "LiteLLM_SpendLogs" WHERE "user" = $1', "u1") + transaction.reads.assert_awaited_once_with('SELECT spend FROM "LiteLLM_SpendLogs" WHERE "user" = $1', "u1") assert prisma.read_db.tx.call_args.kwargs["timeout"] > timedelta(seconds=30), "Prisma must outlive the Postgres cancel" +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_waits_for_a_pooled_connection_as_long_as_a_plain_read() -> None: + """A transaction start must not give up on the pool before a plain ``query_raw`` would, or a busy pool turns + 25 concurrent spend reads into Prisma P2028 failures that the plain path served fine.""" + from datetime import timedelta + + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_reads + + prisma: Final = MagicMock() + prisma.read_db.pool_timeout = timedelta(seconds=60) + statement_timeout_reads(prisma, AsyncMock(return_value=[])) + + await query_raw_within_statement_timeout(prisma, "SELECT 1") + + assert prisma.read_db.tx.call_args.kwargs["max_wait"] == timedelta(seconds=60) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("session_statement_timeout_ms", "expected_statements"), + [ + (0, ("SET LOCAL statement_timeout = 30000",)), + (60_000, ("SET LOCAL statement_timeout = 30000",)), + (30_000, ()), + (500, ()), + ], + ids=["unlimited", "looser", "equal", "stricter"], +) +async def test_query_raw_within_statement_timeout_only_lowers_the_session_statement_timeout( + session_statement_timeout_ms: int, expected_statements: tuple[str, ...] +) -> None: + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_reads + + prisma: Final = MagicMock() + transaction: Final = statement_timeout_reads(prisma, AsyncMock(return_value=[]), session_statement_timeout_ms) + + await query_raw_within_statement_timeout(prisma, "SELECT 1") + + assert transaction.statements == expected_statements + + +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_reports_the_stricter_limit_the_read_ran_under() -> None: + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads + + prisma: Final = MagicMock() + statement_timeout_reads(prisma, AsyncMock(side_effect=statement_timeout_error()), statement_timeout_ms=500) + + with pytest.raises(HTTPException) as exc_info: + await query_raw_within_statement_timeout(prisma, "SELECT 1", what="Spend per user") + + assert exc_info.value.detail == {"error": "Spend per user took longer than 0.5s. Pick a shorter date range"} + + @pytest.mark.asyncio async def test_query_raw_within_statement_timeout_reports_a_cancelled_read_as_503() -> None: from litellm.proxy.utils import query_raw_within_statement_timeout