diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 6c4e2bbfdcf..a7f18769461 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -398,12 +398,26 @@ class PrismaWrapper: self._retirement_tasks = self._retirement_tasks.difference((retirement_task,)) async def connect(self, timeout: int | timedelta | None = None) -> None: - if timeout is None: - await self._original_prisma.connect() - else: - await self._original_prisma.connect(timeout) + await self._connect_prisma(self._original_prisma, timeout) self._active_drain_tracker = self._instrument_prisma_client(self._original_prisma) + @staticmethod + async def _connect_prisma(prisma_client: "Prisma", timeout: int | timedelta | None = None) -> None: + new_connection: Final = prisma_client.is_connected() is False + if new_connection: + from litellm.proxy.db.prisma_engine import initialize_query_engine + + initialize_query_engine(prisma_client) + try: + if timeout is None: + await prisma_client.connect() + else: + await prisma_client.connect(timeout) + except (Exception, asyncio.CancelledError): + if new_connection: + await prisma_client.disconnect() + raise + @staticmethod async def _kill_engine_process(pid: int) -> None: """Force-kill the engine subprocess to prevent DB connection pool leaks. @@ -688,7 +702,7 @@ class PrismaWrapper: else: self._original_prisma = Prisma() - await self._original_prisma.connect() + await self._connect_prisma(self._original_prisma) self._active_drain_tracker = self._instrument_prisma_client(self._original_prisma) self._engine_generation += 1 @@ -719,7 +733,7 @@ class PrismaWrapper: if old_drain_tracker is None: old_drain_tracker = self._instrument_prisma_client(old_prisma) try: - await replacement_prisma.connect() + await self._connect_prisma(replacement_prisma) except (Exception, asyncio.CancelledError): self._schedule_engine_retirement(self._get_engine_pid(replacement_prisma), None) raise diff --git a/litellm/proxy/db/prisma_engine.py b/litellm/proxy/db/prisma_engine.py new file mode 100644 index 00000000000..212988579fc --- /dev/null +++ b/litellm/proxy/db/prisma_engine.py @@ -0,0 +1,136 @@ +import asyncio +import os +import signal +import subprocess +import sys +from collections.abc import Sequence +from datetime import timedelta +from pathlib import Path +from threading import Thread +from typing import TYPE_CHECKING, BinaryIO, Final, Literal + +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +if TYPE_CHECKING: + from prisma import Prisma + from prisma.types import DatasourceOverride + + +class _AddressFields(BaseModel): + model_config = ConfigDict(frozen=True) + ip: Literal["127.0.0.1"] + port: int = Field(gt=0, le=65535) + message: str + + +class _ListeningEvent(BaseModel): + model_config = ConfigDict(frozen=True) + target: Literal["query_engine::server"] + fields: _AddressFields + + +def _listening_url(line: bytes) -> str | None: + try: + event: Final = _ListeningEvent.model_validate_json(line) + except ValidationError: + return None + if not event.fields.message.startswith("Started query engine http server on "): + return None + return f"http://{event.fields.ip}:{event.fields.port}" + + +def _publish_address(ready: asyncio.Future[str | None], address: str | None) -> None: + if not ready.done(): + ready.set_result(address) + + +def _forward_output(stream: BinaryIO, loop: asyncio.AbstractEventLoop, ready: asyncio.Future[str | None]) -> None: + def publish(address: str | None) -> None: + try: + loop.call_soon_threadsafe(_publish_address, ready, address) + except RuntimeError: + pass + + try: + with stream: + for line in stream: + sys.stdout.write(line.decode("utf-8", errors="replace")) + sys.stdout.flush() + if (address := _listening_url(line)) is not None: + publish(address) + finally: + publish(None) + + +def _unblock_signals() -> None: + signal.pthread_sigmask(signal.SIG_UNBLOCK, (signal.SIGINT, signal.SIGTERM)) + + +def initialize_query_engine(client: "Prisma") -> None: + from prisma._constants import DEFAULT_CONNECT_TIMEOUT + from prisma.builder import dumps + from prisma.client import PACKAGED_SCHEMA_PATH + from prisma.engine import errors + from prisma.engine.query import QueryEngine + from prisma.utils import ( + DEBUG, + _env_bool, # pyright: ignore[reportPrivateUsage] # Preserve Prisma's playground flag semantics + ) + + class PortOwningQueryEngine(QueryEngine): + async def spawn( + self, + file: Path, + timeout: timedelta = DEFAULT_CONNECT_TIMEOUT, + datasources: Sequence["DatasourceOverride"] | None = None, + ) -> None: + playground: Final = _env_bool("__PRISMA_PY_PLAYGROUND") + environment: Final = { + **os.environ, + "PRISMA_DML_PATH": str(self.dml_path.absolute()), + "RUST_LOG": "info" if DEBUG or playground else "error,query_engine::server=info", + "RUST_LOG_FORMAT": "json", + "PRISMA_CLIENT_ENGINE_TYPE": "binary", + "PRISMA_ENGINE_PROTOCOL": "graphql", + **({"OVERWRITE_DATASOURCES": dumps(datasources)} if datasources is not None else {}), + **({"LOG_QUERIES": "y"} if self._log_queries else {}), + } + self.process = subprocess.Popen( + ( + str(file.absolute()), + "-p", + "0", + "--enable-metrics", + "--enable-raw-queries", + *(("--enable-playground",) if playground else ()), + ), + env=environment, + stdout=subprocess.PIPE, + stderr=sys.stderr, + preexec_fn=None if sys.platform == "win32" else _unblock_signals, + ) + stream: Final = self.process.stdout + if stream is None: + raise errors.EngineConnectionError("Query engine output pipe is unavailable") + loop: Final = asyncio.get_running_loop() + ready: Final[asyncio.Future[str | None]] = loop.create_future() + Thread(target=_forward_output, args=(stream, loop, ready), daemon=True).start() + try: + async with asyncio.timeout(timeout.total_seconds()): + address: Final = await ready + if address is None: + raise errors.EngineConnectionError( + "Query engine exited before publishing its listening address" + ) + self.url = address + await self.request("GET", "/status") + except TimeoutError as error: + raise errors.EngineConnectionError( + "Timed out waiting for the query engine's listening address" + ) from error + + client._engine = PortOwningQueryEngine( # pyright: ignore[reportPrivateUsage] # Prisma exposes engine injection only through this setter # rebind-ok: install the lifecycle-owned transport + dml_path=PACKAGED_SCHEMA_PATH, + log_queries=client._log_queries, # pyright: ignore[reportPrivateUsage] # Preserve the generated client's query logging + **client._http_config, # pyright: ignore[reportPrivateUsage] # Preserve the generated client's HTTP options + ) diff --git a/tests/integration/database/test_prisma_engine_port.py b/tests/integration/database/test_prisma_engine_port.py new file mode 100644 index 00000000000..20408c94782 --- /dev/null +++ b/tests/integration/database/test_prisma_engine_port.py @@ -0,0 +1,99 @@ +import asyncio +import os +import socket +from datetime import timedelta +from typing import Final +from urllib.parse import urlsplit + +import pytest +from prisma import Prisma +from prisma.engine.errors import AlreadyConnectedError, EngineConnectionError + +from litellm.caching.caching import DualCache +from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.utils import PrismaClient, ProxyLogging +from tests.integration._support.database import scratch_database + + +@pytest.mark.asyncio +@pytest.mark.parametrize("occupied", (False, True), ids=("normal", "occupied-port")) +async def test_prisma_connect_owns_its_listening_port(occupied: bool, monkeypatch: pytest.MonkeyPatch) -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as competitor: + competitor.bind(("127.0.0.1", 0)) # test-quality-ok: keep the socket bound to force Prisma's port collision + port: Final = competitor.getsockname()[1] + if occupied: + monkeypatch.setattr("prisma.engine.utils.get_open_port", lambda: port) + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(DualCache())) + try: + await client.db.connect(timeout=timedelta(seconds=5)) + with pytest.raises(AlreadyConnectedError): + await client.db.connect(timeout=timedelta(seconds=5)) + assert await client.db.query_raw("SELECT 1 AS result") == [{"result": 1}] + assert client.db.is_connected() is True + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_concurrent_prisma_engines_query_their_own_databases() -> None: + with scratch_database("engine_first") as first, scratch_database("engine_second") as second: + clients: Final = tuple(PrismaWrapper(Prisma(datasource={"url": url})) for url in (first, second)) + try: + await asyncio.gather(*(client.connect() for client in clients)) + results: Final = await asyncio.gather( + *(client.query_raw("SELECT current_database() AS name") for client in clients) + ) + assert results == [[{"name": urlsplit(url).path.lstrip("/")}] for url in (first, second)] + finally: + await asyncio.gather(*(client.disconnect() for client in clients)) + + +@pytest.mark.asyncio +async def test_prisma_startup_database_failure_closes_the_engine(monkeypatch: pytest.MonkeyPatch) -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as unavailable: + unavailable.bind( # test-quality-ok: socket remains bound throughout the connection attempt + ("127.0.0.1", 0) + ) + url: Final = ( + f"postgresql://invalid:invalid@127.0.0.1:{unavailable.getsockname()[1]}/unavailable?connect_timeout=1" + ) + original: Final = Prisma(datasource={"url": url}) + client: Final = PrismaWrapper(original) + connecting: Final = asyncio.create_task(client.connect(timeout=timedelta(seconds=5))) + try: + await asyncio.sleep(0) + engine: Final = original._engine + with pytest.raises(EngineConnectionError): + await connecting + assert original.is_connected() is False + assert engine.process is None + assert engine.session.should_close() is False + monkeypatch.setattr(original, "_datasource", {"url": os.environ["DATABASE_URL"]}) + await client.connect(timeout=timedelta(seconds=5)) + assert await client.query_raw("SELECT 1 AS result") == [{"result": 1}] + assert original.is_connected() is True + finally: + await client.disconnect() + + +@pytest.mark.asyncio +async def test_cancelled_prisma_startup_stops_its_engine() -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as stalled_database: + stalled_database.bind(("127.0.0.1", 0)) # test-quality-ok: socket remains bound during the stalled startup + stalled_database.listen() + url: Final = f"postgresql://invalid:invalid@127.0.0.1:{stalled_database.getsockname()[1]}/stalled" + original: Final = Prisma(datasource={"url": url}) + client: Final = PrismaWrapper(original) + connecting: Final = asyncio.create_task(client.connect()) + try: + await asyncio.sleep(0) + process: Final = original._engine.process + assert process is not None + connecting.cancel() + with pytest.raises(asyncio.CancelledError): + await connecting + assert process.poll() is not None + assert original.is_connected() is False + finally: + await client.disconnect() diff --git a/tests/unit/proxy/db/test_prisma_engine.py b/tests/unit/proxy/db/test_prisma_engine.py new file mode 100644 index 00000000000..84ee45b3cb9 --- /dev/null +++ b/tests/unit/proxy/db/test_prisma_engine.py @@ -0,0 +1,72 @@ +import asyncio +from typing import Final, NoReturn + +import pytest +from prisma import Prisma +from prisma.engine.errors import EngineConnectionError +from prisma.engine.query import QueryEngine + +from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.db.prisma_engine import _listening_url + + +@pytest.mark.parametrize( + ("line", "expected"), + ( + ( + b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"31415",' + b'"message":"Started query engine http server on http://127.0.0.1:31415"}}', + "http://127.0.0.1:31415", + ), + (b"", None), + (b"query-engine version", None), + (b'{"target":"query_engine","fields":{"message":"PANIC","reason":"Address already in use"}}', None), + ( + b'{"target":"other","fields":{"ip":"127.0.0.1","port":"31415",' + b'"message":"Started query engine http server on http://127.0.0.1:31415"}}', + None, + ), + ( + b'{"target":"query_engine::server","fields":{"ip":"0.0.0.0","port":"31415",' + b'"message":"Started query engine http server on http://0.0.0.0:31415"}}', + None, + ), + ( + b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"0",' + b'"message":"Started query engine http server on http://127.0.0.1:0"}}', + None, + ), + ( + b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"65536",' + b'"message":"Started query engine http server on http://127.0.0.1:65536"}}', + None, + ), + ( + b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"31415",' + b'"message":"Other server event"}}', + None, + ), + ), +) +def test_listening_url_only_accepts_an_engine_listening_event(line: bytes, expected: str | None) -> None: + assert _listening_url(line) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fault", (EngineConnectionError("engine startup failed"), asyncio.CancelledError())) +async def test_failed_engine_startup_clears_connection_state( + fault: EngineConnectionError | asyncio.CancelledError, monkeypatch: pytest.MonkeyPatch +) -> None: + def unavailable_engine(_: QueryEngine) -> NoReturn: + raise fault + + monkeypatch.setattr(QueryEngine, "_ensure_file", unavailable_engine) + original: Final = Prisma(datasource={"url": "postgresql://unused@localhost/unused"}) + client: Final = PrismaWrapper(original) + try: + with pytest.raises(type(fault)) as caught: + await client.connect() + assert caught.value is fault + assert original.is_connected() is False + finally: + await client.disconnect() diff --git a/tests/unit/proxy/db/test_routing_prisma_wrapper.py b/tests/unit/proxy/db/test_routing_prisma_wrapper.py index 73225acd2c8..285db852ef9 100644 --- a/tests/unit/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/unit/proxy/db/test_routing_prisma_wrapper.py @@ -574,6 +574,9 @@ async def test_writer_recreate_passes_http_client_through(monkeypatch): def __init__(self, **kwargs): captured_kwargs.update(kwargs) + def is_connected(self): + return True + async def connect(self): return None @@ -875,6 +878,9 @@ async def test_reader_recreate_uses_datasource_override(monkeypatch): def __init__(self, **kwargs): captured_kwargs.update(kwargs) + def is_connected(self): + return True + async def connect(self): return None @@ -912,6 +918,9 @@ async def test_writer_recreate_does_not_use_datasource(monkeypatch): def __init__(self, **kwargs): captured_kwargs.update(kwargs) + def is_connected(self): + return True + async def connect(self): return None diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 05667629305..27a5734eb45 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8298,6 +8298,8 @@ class _ReadOnlyReplicaPrisma: def __init__(self, journal: _PrismaJournal, role: str) -> None: self.journal = journal self.role = role + self._log_queries: Final = False + self._http_config: Final[dict[str, object]] = {} def is_connected(self) -> bool: return False diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 0bf7d199205..5fa86936184 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -2713,7 +2713,7 @@ async def test_proxy_server_prisma_setup_invalid_db(monkeypatch): Think 2-3 times before editing / deleting this test, it's important for PROD """ - import httpx + from prisma.engine.errors import EngineConnectionError from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging @@ -2724,7 +2724,7 @@ async def test_proxy_server_prisma_setup_invalid_db(monkeypatch): monkeypatch.setenv("DATABASE_URL", invalid_db_url) - with pytest.raises(httpx.ConnectError): + with pytest.raises(EngineConnectionError): await ProxyStartupEvent._setup_prisma_client( database_url=invalid_db_url, proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),