mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
685c3ec6a2
Cancelling a task mid-request closed the query engine socket, the engine handed the still-busy pool connection back, the next transaction start queued its BEGIN behind the orphaned statement, and when that waiter was cancelled too the engine panicked on its dropped reply channel (interactive_transactions/actors.rs) and exited under the shutdown spend flush, losing daily user and team spend for requests that had returned 200. Every engine request now runs in its own task behind asyncio.shield so only the caller is cancelled, and a transaction that opens after its caller left is rolled back. Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
574 lines
21 KiB
Python
574 lines
21 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import signal
|
|
import sys
|
|
import urllib.parse
|
|
from typing import Any, Final, cast
|
|
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
|
|
from litellm.proxy.db.prisma_client import (
|
|
PrismaManager,
|
|
PrismaWrapper,
|
|
_PrismaDrainTracker,
|
|
_TrackedPrismaEngine,
|
|
should_update_prisma_schema,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_prisma_binary():
|
|
"""Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests."""
|
|
mock_module = MagicMock()
|
|
with patch.dict(sys.modules, {"prisma": mock_module}):
|
|
yield mock_module
|
|
|
|
|
|
def test_should_update_prisma_schema(monkeypatch):
|
|
# CASE 1: Environment variable behavior
|
|
# When DISABLE_SCHEMA_UPDATE is not set -> should update
|
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
|
|
assert should_update_prisma_schema() == True
|
|
|
|
# When DISABLE_SCHEMA_UPDATE="true" -> should not update
|
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "true")
|
|
assert should_update_prisma_schema() == False
|
|
|
|
# When DISABLE_SCHEMA_UPDATE="false" -> should update
|
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "false")
|
|
assert should_update_prisma_schema() == True
|
|
|
|
# CASE 2: Explicit parameter behavior (overrides env var)
|
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
|
|
assert should_update_prisma_schema(True) == False # Param True -> should not update
|
|
|
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None) # Set env var opposite to param
|
|
assert should_update_prisma_schema(False) == True # Param False -> should update
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recreate_prisma_client_successful_disconnect():
|
|
"""
|
|
Test that recreate_prisma_client works normally when disconnect succeeds.
|
|
"""
|
|
# Mock the original prisma client
|
|
mock_prisma = AsyncMock()
|
|
|
|
# Create a mock PrismaWrapper instance
|
|
wrapper = Mock()
|
|
wrapper._original_prisma = mock_prisma
|
|
|
|
# Configure disconnect to succeed
|
|
mock_prisma.disconnect.return_value = None
|
|
|
|
# Mock the entire recreate_prisma_client method to avoid import issues
|
|
async def mock_recreate_prisma_client(new_db_url: str, http_client=None):
|
|
try:
|
|
await mock_prisma.disconnect()
|
|
except Exception:
|
|
pass
|
|
|
|
mock_new_prisma = AsyncMock()
|
|
wrapper._original_prisma = mock_new_prisma
|
|
await mock_new_prisma.connect()
|
|
|
|
# Assign the mock method to the wrapper
|
|
wrapper.recreate_prisma_client = mock_recreate_prisma_client
|
|
|
|
# Call the method
|
|
await wrapper.recreate_prisma_client("postgresql://new:new@localhost:5432/new")
|
|
|
|
# Verify that disconnect was called
|
|
mock_prisma.disconnect.assert_called_once()
|
|
|
|
# Verify that the new client replaced the original
|
|
assert wrapper._original_prisma != mock_prisma
|
|
assert hasattr(wrapper._original_prisma, "connect")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recreate_prisma_client_kills_old_engine_on_disconnect_failure(
|
|
mock_prisma_binary,
|
|
):
|
|
"""When disconnect() fails, recreate_prisma_client must SIGTERM/SIGKILL the old engine PID."""
|
|
mock_prisma = AsyncMock()
|
|
mock_prisma.disconnect.side_effect = Exception("engine hung")
|
|
mock_prisma.is_connected = MagicMock(return_value=True)
|
|
|
|
# Simulate engine subprocess with a known PID
|
|
mock_engine = MagicMock()
|
|
mock_engine.process.pid = 12345
|
|
mock_prisma._engine = mock_engine
|
|
|
|
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
|
|
|
|
# Configure the mock Prisma constructor
|
|
mock_new_prisma = AsyncMock()
|
|
mock_prisma_binary.Prisma.return_value = mock_new_prisma
|
|
|
|
with (
|
|
patch("os.kill") as mock_kill,
|
|
patch("asyncio.sleep", new_callable=AsyncMock),
|
|
):
|
|
await wrapper.recreate_prisma_client("postgresql://new")
|
|
|
|
# Verify old engine was killed
|
|
mock_kill.assert_any_call(12345, signal.SIGTERM)
|
|
# Verify new client was created and connected
|
|
mock_new_prisma.connect.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recreate_prisma_client_skips_kill_on_successful_disconnect(
|
|
mock_prisma_binary,
|
|
):
|
|
"""When disconnect() succeeds, no kill should be attempted."""
|
|
mock_prisma = AsyncMock()
|
|
mock_prisma.is_connected = MagicMock(return_value=True)
|
|
mock_prisma.disconnect.return_value = None
|
|
|
|
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
|
|
|
|
mock_new_prisma = AsyncMock()
|
|
mock_prisma_binary.Prisma.return_value = mock_new_prisma
|
|
|
|
with patch("os.kill") as mock_kill:
|
|
await wrapper.recreate_prisma_client("postgresql://new")
|
|
|
|
mock_kill.assert_not_called()
|
|
mock_new_prisma.connect.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recreate_prisma_client_handles_missing_engine_pid(
|
|
mock_prisma_binary,
|
|
):
|
|
"""When engine PID is unavailable (no _engine attr), kill is skipped gracefully."""
|
|
mock_prisma = AsyncMock()
|
|
mock_prisma.is_connected = MagicMock(return_value=True)
|
|
mock_prisma.disconnect.side_effect = Exception("engine hung")
|
|
mock_prisma._engine = None # No engine subprocess
|
|
|
|
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
|
|
|
|
mock_new_prisma = AsyncMock()
|
|
mock_prisma_binary.Prisma.return_value = mock_new_prisma
|
|
|
|
with (
|
|
patch("os.kill") as mock_kill,
|
|
patch("asyncio.sleep", new_callable=AsyncMock),
|
|
):
|
|
await wrapper.recreate_prisma_client("postgresql://new")
|
|
|
|
mock_kill.assert_not_called() # PID was 0, kill skipped
|
|
mock_new_prisma.connect.assert_awaited_once()
|
|
|
|
|
|
def test_get_engine_pid_returns_zero_for_disconnected_client(disconnected_prisma):
|
|
"""A disconnected client must read as "no engine" instead of raising,
|
|
otherwise the reconnect path can never recover."""
|
|
wrapper = PrismaWrapper(
|
|
original_prisma=disconnected_prisma, iam_token_db_auth=False
|
|
)
|
|
|
|
assert wrapper._get_engine_pid() == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recreate_prisma_client_recovers_from_disconnected_client(
|
|
mock_prisma_binary, disconnected_prisma
|
|
):
|
|
"""recreate_prisma_client must still build a replacement client when the
|
|
current one is disconnected."""
|
|
wrapper = PrismaWrapper(
|
|
original_prisma=disconnected_prisma, iam_token_db_auth=False
|
|
)
|
|
|
|
mock_new_prisma = AsyncMock()
|
|
mock_prisma_binary.Prisma.return_value = mock_new_prisma
|
|
|
|
with patch("os.kill") as mock_kill:
|
|
result = await wrapper.recreate_prisma_client("postgresql://new")
|
|
|
|
assert result is True
|
|
mock_kill.assert_not_called()
|
|
assert wrapper._original_prisma is mock_new_prisma
|
|
mock_new_prisma.connect.assert_awaited_once()
|
|
|
|
|
|
DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"]
|
|
|
|
|
|
def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url):
|
|
"""`prisma db push` bypasses litellm-proxy-extras, so it needs its own call
|
|
into the opt-in REPLICA IDENTITY FULL step."""
|
|
from litellm.proxy.db.prisma_client import PrismaManager
|
|
from litellm_proxy_extras.replica_identity import REPLICA_IDENTITY_FULL_ENV_VAR
|
|
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
|
|
|
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
|
|
applied = []
|
|
monkeypatch.setattr(
|
|
ProxyExtrasDBManager,
|
|
"apply_replica_identity_full_if_requested",
|
|
staticmethod(lambda: applied.append(True)),
|
|
)
|
|
|
|
assert PrismaManager.setup_database(use_migrate=False) is True
|
|
|
|
assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
|
|
assert applied == [True]
|
|
|
|
|
|
def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
|
|
"""A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the
|
|
primary key back to ("request_id"), which Postgres rejects; the guard must
|
|
fail fast with guidance instead of running the push."""
|
|
from litellm.proxy.db.prisma_client import PrismaManager
|
|
from litellm_proxy_extras.utils import (
|
|
PARTITIONED_SPEND_LOGS_PUSH_ERROR,
|
|
ProxyExtrasDBManager,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
|
|
)
|
|
with pytest.raises(RuntimeError) as err:
|
|
PrismaManager.setup_database(use_migrate=False)
|
|
|
|
assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
|
|
assert fake_prisma_cli.calls == []
|
|
|
|
|
|
def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
|
|
from litellm.proxy.db.prisma_client import PrismaManager
|
|
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
|
|
|
monkeypatch.setattr(
|
|
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False)
|
|
)
|
|
assert PrismaManager.setup_database(use_migrate=False) is True
|
|
|
|
assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
|
|
|
|
|
|
def _entra_jwt(expires_in_seconds: int) -> str:
|
|
"""A JWT shaped like a real Entra access token, expiring ``expires_in_seconds`` from now."""
|
|
import base64
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
exp = int((datetime.now(tz=timezone.utc) + timedelta(seconds=expires_in_seconds)).timestamp())
|
|
payload = base64.urlsafe_b64encode(json.dumps({"exp": exp}).encode()).rstrip(b"=").decode()
|
|
return f"aGVhZGVy.{payload}.c2ln"
|
|
|
|
|
|
@pytest.fixture
|
|
def azure_env(monkeypatch, unset_database_url):
|
|
monkeypatch.setenv("DATABASE_HOST", "pg.postgres.database.azure.com")
|
|
monkeypatch.setenv("DATABASE_PORT", "5432")
|
|
monkeypatch.setenv("DATABASE_USER", "litellm@contoso.onmicrosoft.com")
|
|
monkeypatch.setenv("DATABASE_NAME", "litellm_db")
|
|
|
|
|
|
def _azure_wrapper(token: str, **kwargs):
|
|
from litellm.proxy.db.token_auth import AzureEntraTokenAuth
|
|
|
|
return PrismaWrapper(
|
|
original_prisma=MagicMock(),
|
|
token_auth=AzureEntraTokenAuth(token_provider=lambda: token),
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def test_azure_entra_mint_writes_an_encoded_url_into_the_db_url_env_var(azure_env):
|
|
"""The UPN user and the JWT both have to survive being embedded in a URL."""
|
|
token = _entra_jwt(3600)
|
|
wrapper = _azure_wrapper(token)
|
|
|
|
db_url = wrapper.get_rds_iam_token()
|
|
|
|
assert db_url == (
|
|
f"postgresql://litellm%40contoso.onmicrosoft.com:{urllib.parse.quote(token, safe='')}"
|
|
"@pg.postgres.database.azure.com:5432/litellm_db"
|
|
)
|
|
assert os.environ["DATABASE_URL"] == db_url
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("previous_query", "expected_query"),
|
|
[
|
|
("max_idle_connection_lifetime=60", {"max_idle_connection_lifetime": ["60"]}),
|
|
(
|
|
"connection_limit=20&pgbouncer=true&max_idle_connection_lifetime=45",
|
|
{"connection_limit": ["20"], "pgbouncer": ["true"], "max_idle_connection_lifetime": ["45"]},
|
|
),
|
|
(
|
|
"sslmode=require&sslcert=/certs/root.pem&sslaccept=strict&schema=tenant",
|
|
{"sslmode": ["require"], "sslcert": ["/certs/root.pem"], "sslaccept": ["strict"]},
|
|
),
|
|
],
|
|
)
|
|
def test_token_refresh_keeps_the_connection_params_of_the_url_it_replaces(
|
|
azure_env, monkeypatch, previous_query, expected_query
|
|
):
|
|
old_token = _entra_jwt(60)
|
|
monkeypatch.setenv(
|
|
"DATABASE_URL",
|
|
f"postgresql://litellm%40contoso.onmicrosoft.com:{urllib.parse.quote(old_token, safe='')}"
|
|
f"@pg.postgres.database.azure.com:5432/litellm_db?{previous_query}",
|
|
)
|
|
new_token = _entra_jwt(3600)
|
|
|
|
db_url = _azure_wrapper(new_token).get_rds_iam_token()
|
|
|
|
assert db_url is not None
|
|
assert os.environ["DATABASE_URL"] == db_url
|
|
assert urllib.parse.quote(new_token, safe="") in db_url
|
|
assert urllib.parse.parse_qs(urllib.parse.urlsplit(db_url).query) == expected_query
|
|
|
|
|
|
def test_token_refresh_keeps_the_reader_url_params_separate_from_the_writer(azure_env, monkeypatch):
|
|
from litellm.proxy.db.token_auth import IAMEndpoint
|
|
|
|
monkeypatch.setenv("DATABASE_URL", "postgresql://w:t@pg:5432/litellm_db?max_idle_connection_lifetime=45")
|
|
monkeypatch.setenv(
|
|
"DATABASE_URL_READ_REPLICA", "postgresql://r:t@replica:5432/litellm_db?max_idle_connection_lifetime=60"
|
|
)
|
|
reader = _azure_wrapper(
|
|
_entra_jwt(3600),
|
|
db_url_env_var="DATABASE_URL_READ_REPLICA",
|
|
iam_endpoint=IAMEndpoint(host="replica", port="5432", user="r", name="litellm_db", schema=None),
|
|
)
|
|
|
|
reader_url = reader.get_rds_iam_token()
|
|
|
|
assert reader_url is not None
|
|
assert reader_url.startswith("postgresql://r:")
|
|
assert urllib.parse.parse_qs(urllib.parse.urlsplit(reader_url).query) == {"max_idle_connection_lifetime": ["60"]}
|
|
assert os.environ["DATABASE_URL"].endswith("?max_idle_connection_lifetime=45")
|
|
|
|
|
|
def test_azure_entra_refresh_is_scheduled_off_the_jwt_expiry(azure_env):
|
|
"""Without reading `exp` this falls back to a fixed 600s interval, which silently
|
|
outlives a token and breaks every reconnect after it lapses (issue #29661)."""
|
|
wrapper = _azure_wrapper(_entra_jwt(3600))
|
|
wrapper.get_rds_iam_token()
|
|
|
|
seconds = wrapper._calculate_seconds_until_refresh()
|
|
|
|
expected = 3600 - PrismaWrapper.TOKEN_REFRESH_BUFFER_SECONDS
|
|
assert seconds != PrismaWrapper.FALLBACK_REFRESH_INTERVAL_SECONDS
|
|
assert expected - 5 <= seconds <= expected
|
|
|
|
|
|
def test_a_token_whose_expiry_never_advances_cannot_spin_the_refresh_loop(azure_env):
|
|
"""azure-identity hands back its cached token when a renewal attempt fails inside its
|
|
own window, so a transient Entra or IMDS problem in the last 3 minutes of a token
|
|
yields a successful refresh whose `exp` has not moved. With no floor on the sleep the
|
|
loop then re-mints and recreates the query engine on every pass, with nothing in
|
|
between, for as long as Entra stays sick."""
|
|
wrapper = _azure_wrapper(_entra_jwt(60))
|
|
wrapper.get_rds_iam_token()
|
|
first = wrapper._calculate_seconds_until_refresh()
|
|
|
|
wrapper.get_rds_iam_token()
|
|
second = wrapper._calculate_seconds_until_refresh()
|
|
|
|
assert first == second == PrismaWrapper.TOKEN_REFRESH_MIN_SLEEP_SECONDS
|
|
|
|
|
|
def test_azure_entra_token_expiry_is_detected(azure_env):
|
|
wrapper = _azure_wrapper(_entra_jwt(3600))
|
|
fresh_url = wrapper.get_rds_iam_token()
|
|
expired_url = _azure_wrapper(_entra_jwt(-1)).get_rds_iam_token()
|
|
|
|
assert wrapper.is_token_expired(fresh_url) is False
|
|
assert wrapper.is_token_expired(expired_url) is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_entra_strategy_starts_the_refresh_task(azure_env):
|
|
"""The refresh loop is gated on the legacy boolean, so an Azure strategy has to
|
|
get past that gate; a password-auth wrapper still must not start a task."""
|
|
wrapper = _azure_wrapper(_entra_jwt(3600))
|
|
wrapper.get_rds_iam_token()
|
|
password_wrapper = PrismaWrapper(original_prisma=MagicMock())
|
|
|
|
await wrapper.start_token_refresh_task()
|
|
await password_wrapper.start_token_refresh_task()
|
|
try:
|
|
assert wrapper._token_refresh_task is not None
|
|
assert not wrapper._token_refresh_task.done()
|
|
assert password_wrapper._token_refresh_task is None
|
|
finally:
|
|
await wrapper.stop_token_refresh_task()
|
|
|
|
|
|
def test_azure_entra_strategy_reads_as_token_auth_enabled(azure_env):
|
|
"""`routing_prisma_wrapper` gates the reader's refresh on this flag, so an Azure
|
|
reader has to answer True to it."""
|
|
wrapper = _azure_wrapper(_entra_jwt(3600))
|
|
|
|
assert wrapper.iam_token_db_auth is True
|
|
assert wrapper.token_label == "Azure Entra token"
|
|
|
|
|
|
def test_the_token_strategy_cannot_be_swapped_after_construction(azure_env):
|
|
"""Assigning the legacy boolean used to replace a configured Entra strategy with the
|
|
RDS one, which points boto at an Azure host."""
|
|
wrapper = _azure_wrapper(_entra_jwt(3600))
|
|
|
|
with pytest.raises(AttributeError):
|
|
wrapper.iam_token_db_auth = True
|
|
|
|
|
|
def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch):
|
|
"""A blank host used to produce `postgresql://:<token>@:5432/`, which fails deep
|
|
inside Prisma instead of at the misconfiguration."""
|
|
monkeypatch.delenv("DATABASE_HOST")
|
|
wrapper = _azure_wrapper(_entra_jwt(3600))
|
|
|
|
with pytest.raises(RuntimeError, match="DATABASE_HOST"):
|
|
wrapper.get_rds_iam_token()
|
|
|
|
|
|
@pytest.mark.timeout(45)
|
|
def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch):
|
|
"""
|
|
A timed-out `db push` used to leave Node and the schema engine writing the schema,
|
|
so the next attempt pushed into a database the abandoned one was still mutating.
|
|
"""
|
|
monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False)
|
|
monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
|
|
|
|
assert PrismaManager.setup_database(use_migrate=False) is True
|
|
assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV]
|
|
assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
|
|
|
|
|
|
def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot(
|
|
fake_prisma_cli, unset_database_url, monkeypatch
|
|
):
|
|
"""
|
|
An ImportError out of setup_database escapes the caller's RuntimeError handler and
|
|
kills boot with a traceback instead of the failed-setup message and exit code.
|
|
"""
|
|
monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
|
|
|
|
assert PrismaManager.setup_database(use_migrate=False) is False
|
|
assert fake_prisma_cli.calls == []
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("run", "outcome"),
|
|
(
|
|
(PrismaManager.build_request_log_indexes, False),
|
|
(PrismaManager.start_request_log_index_build, None),
|
|
),
|
|
ids=("wait-for-the-build", "start-the-build"),
|
|
)
|
|
def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising(monkeypatch, run, outcome):
|
|
"""The migration job exits non-zero and a serving proxy keeps booting when the extras
|
|
package that owns the index build is not installed."""
|
|
monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None)
|
|
|
|
assert run() is outcome
|
|
|
|
|
|
class _SlowEngine:
|
|
"""Engine double whose requests only answer once the test lets them, recording what the engine saw."""
|
|
|
|
def __init__(self) -> None:
|
|
self.gates: dict[str, asyncio.Event] = {} # mutable-ok: one gate per request the test lets through
|
|
self.queries_finished = 0
|
|
self.requests_cancelled = 0
|
|
self.rolled_back: list[str] = [] # mutable-ok: test records
|
|
self.process = object()
|
|
|
|
def gate(self, request: str) -> asyncio.Event:
|
|
return self.gates.setdefault(request, asyncio.Event())
|
|
|
|
async def _answer_once_released(self, request: str) -> None:
|
|
try:
|
|
await self.gate(request).wait()
|
|
except asyncio.CancelledError:
|
|
self.requests_cancelled += 1
|
|
raise
|
|
|
|
async def query(self, content: str, *, tx_id: str | None) -> object:
|
|
await self._answer_once_released(content)
|
|
self.queries_finished += 1
|
|
return {"content": content, "tx_id": tx_id}
|
|
|
|
async def start_transaction(self, *, content: str) -> str:
|
|
await self._answer_once_released("start")
|
|
return "tx-opened-after-caller-left"
|
|
|
|
async def commit_transaction(self, tx_id: str) -> None:
|
|
raise AssertionError(f"unexpected commit of {tx_id}")
|
|
|
|
async def rollback_transaction(self, tx_id: str) -> None:
|
|
self.rolled_back.append(tx_id)
|
|
|
|
|
|
def _tracked(engine: _SlowEngine) -> tuple[_TrackedPrismaEngine, _PrismaDrainTracker]:
|
|
tracker: Final = _PrismaDrainTracker()
|
|
return _TrackedPrismaEngine(cast(Any, engine), tracker), tracker
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelling_a_query_caller_lets_the_engine_finish_the_statement() -> None:
|
|
engine: Final = _SlowEngine()
|
|
tracked, tracker = _tracked(engine)
|
|
caller: Final = asyncio.create_task(tracked.query("x", tx_id=None))
|
|
await asyncio.sleep(0)
|
|
caller.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await caller
|
|
drained: Final = asyncio.create_task(tracker.wait_until_drained())
|
|
await asyncio.sleep(0)
|
|
assert engine.requests_cancelled == 0
|
|
assert not drained.done()
|
|
engine.gate("x").set()
|
|
await asyncio.wait_for(drained, timeout=1)
|
|
assert engine.queries_finished == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transaction_opened_after_its_caller_was_cancelled_is_rolled_back() -> None:
|
|
engine: Final = _SlowEngine()
|
|
tracked, tracker = _tracked(engine)
|
|
caller: Final = asyncio.create_task(tracked.start_transaction(content="{}"))
|
|
await asyncio.sleep(0)
|
|
caller.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await caller
|
|
engine.gate("start").set()
|
|
await asyncio.wait_for(tracker.wait_until_drained(), timeout=1)
|
|
assert engine.requests_cancelled == 0
|
|
assert engine.rolled_back == ["tx-opened-after-caller-left"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain_open() -> None:
|
|
engine: Final = _SlowEngine()
|
|
tracked, tracker = _tracked(engine)
|
|
first: Final = asyncio.create_task(tracked.query("first", tx_id=None))
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
drained: Final = asyncio.create_task(tracker.wait_until_drained())
|
|
await asyncio.sleep(0)
|
|
second: Final = asyncio.create_task(tracked.query("second", tx_id=None))
|
|
engine.gate("first").set()
|
|
await first
|
|
await asyncio.sleep(0)
|
|
assert not drained.done()
|
|
engine.gate("second").set()
|
|
await second
|
|
await asyncio.wait_for(drained, timeout=1)
|
|
assert engine.queries_finished == 2
|