import asyncio import logging import os import sys from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock, call, patch from urllib.parse import parse_qs, urlsplit import pytest # NOTE: do NOT patch sys.modules["prisma"] file-wide via an autouse fixture. # Doing so leaks across pytest-xdist test scheduling: when a worker runs a # routing test, then later runs test_exception_handler.py, the cached MagicMock # attribute references break `isinstance(e, prisma.errors.X)` in # `is_database_transport_error`. The two tests below that actually need to # stub the prisma SDK do so per-test via monkeypatch, which is properly scoped. def _make_wrappers(): from litellm.proxy.db.prisma_client import PrismaWrapper writer_inner = MagicMock(name="writer_prisma") reader_inner = MagicMock(name="reader_prisma") writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False) reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False) return writer, writer_inner, reader, reader_inner class _FakeActions: """Stand-in for a Prisma per-model Actions instance (non-callable, has find_many/create).""" def __init__(self, name: str): self._name = name for method in ( "find_many", "find_unique", "find_first", "count", "group_by", "create", "update", "upsert", "delete", "delete_many", "update_many", ): setattr(self, method, MagicMock(name=f"{name}.{method}")) def _model_actions_mock(name: str) -> _FakeActions: return _FakeActions(name) def test_top_level_query_raw_routes_to_reader(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() routing = RoutingPrismaWrapper(writer=writer, reader=reader) # query_raw should resolve to the reader's underlying client. assert routing.query_raw is reader_inner.query_raw assert routing.query_first is reader_inner.query_first def test_top_level_execute_raw_routes_to_writer(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() routing = RoutingPrismaWrapper(writer=writer, reader=reader) # execute_raw, batch_, tx are write-side and must hit the writer. assert routing.execute_raw is writer_inner.execute_raw assert routing.batch_ is writer_inner.batch_ assert routing.tx is writer_inner.tx def test_per_model_reads_route_to_reader_writes_to_writer(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.litellm_usertable = _model_actions_mock("writer_users") reader_inner.litellm_usertable = _model_actions_mock("reader_users") routing = RoutingPrismaWrapper(writer=writer, reader=reader) actions = routing.litellm_usertable # Reads → reader actions. assert actions.find_many is reader_inner.litellm_usertable.find_many assert actions.find_unique is reader_inner.litellm_usertable.find_unique assert actions.find_first is reader_inner.litellm_usertable.find_first assert actions.count is reader_inner.litellm_usertable.count assert actions.group_by is reader_inner.litellm_usertable.group_by # Writes → writer actions. assert actions.create is writer_inner.litellm_usertable.create assert actions.update is writer_inner.litellm_usertable.update assert actions.upsert is writer_inner.litellm_usertable.upsert assert actions.delete is writer_inner.litellm_usertable.delete assert actions.update_many is writer_inner.litellm_usertable.update_many assert actions.delete_many is writer_inner.litellm_usertable.delete_many def test_writer_pinned_client_bypasses_reader_routing(): """Regression for #38556: read-after-write reconciles must see the writer's just-committed rows, so WriterPinnedClient must resolve reads to the writer even when a read replica is configured.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") routing = RoutingPrismaWrapper(writer=writer, reader=reader) pinned = WriterPinnedClient(routing) assert pinned.db is writer assert pinned.db.litellm_proxymodeltable.find_many is writer_inner.litellm_proxymodeltable.find_many def test_writer_pinned_client_passes_through_single_db(): from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient writer, _, _, _ = _make_wrappers() assert WriterPinnedClient(writer).db is writer def test_writer_pinned_client_yields_to_routed_reads_when_writer_down(): """The pin must not break reader-only degraded mode: a proxy that starts during a primary outage still loads DB-backed models from the replica, so while the writer is degraded the pin resolves to the routed wrapper.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._writer_unavailable = True pinned = WriterPinnedClient(routing) assert pinned.db is routing assert pinned.db.litellm_proxymodeltable.find_many is reader_inner.litellm_proxymodeltable.find_many def test_writer_wrapper_keeps_raw_sql_on_the_writer_while_writer_flagged_down(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, writer_wrapper writer, writer_inner, reader, reader_inner = _make_wrappers() routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._writer_unavailable = True assert writer_wrapper(routing).query_raw is writer_inner.query_raw assert writer_wrapper(routing).query_raw is not reader_inner.query_raw assert writer_wrapper(writer) is writer @pytest.mark.asyncio async def test_connect_invokes_both_clients(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.connect = AsyncMock() reader_inner.connect = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) await routing.connect() writer_inner.connect.assert_awaited_once() reader_inner.connect.assert_awaited_once() @pytest.mark.asyncio async def test_connect_logs_writer_and_reader_success(caplog): """Successful startup emits a positive INFO confirmation for both writer and reader so operators can verify connectivity without inspecting the URL in logs.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.connect = AsyncMock() reader_inner.connect = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): await routing.connect() messages = [r.getMessage() for r in caplog.records] assert "[writer] DB connected" in messages assert "[reader] DB connected" in messages @pytest.mark.asyncio async def test_disconnect_continues_when_one_side_fails(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.disconnect = AsyncMock(side_effect=RuntimeError("writer down")) reader_inner.disconnect = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) with pytest.raises(RuntimeError, match="writer down"): await routing.disconnect() # Reader still attempted even though writer raised. reader_inner.disconnect.assert_awaited_once() def test_is_connected_reflects_writer_only(): """is_connected() must NOT depend on reader health — a healthy writer with a degraded reader should report True so that PrismaClient.connect()'s health check does not re-trigger a writer reconnect (which only fixes writer-side problems and would loop indefinitely).""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() routing = RoutingPrismaWrapper(writer=writer, reader=reader) writer_inner.is_connected = MagicMock(return_value=True) reader_inner.is_connected = MagicMock(return_value=True) assert routing.is_connected() is True # Reader down → still True (reader degradation is tracked separately). reader_inner.is_connected = MagicMock(return_value=False) assert routing.is_connected() is True # Writer down → False. writer_inner.is_connected = MagicMock(return_value=False) assert routing.is_connected() is False def test_token_refresh_delegates_to_both_writer_and_reader(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.start_token_refresh_task = AsyncMock() writer.stop_token_refresh_task = AsyncMock() reader = MagicMock() reader.start_token_refresh_task = AsyncMock() reader.stop_token_refresh_task = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) asyncio.run(routing.start_token_refresh_task()) asyncio.run(routing.stop_token_refresh_task()) # Both wrappers get start/stop — each manages its own IAM token. When # IAM is disabled on a wrapper its task body is a no-op. writer.start_token_refresh_task.assert_awaited_once() writer.stop_token_refresh_task.assert_awaited_once() reader.start_token_refresh_task.assert_awaited_once() reader.stop_token_refresh_task.assert_awaited_once() def test_routed_actions_falls_back_to_writer_for_unknown_methods(): from litellm.proxy.db.routing_prisma_wrapper import _RoutedActions writer_actions = _model_actions_mock("writer") writer_actions.some_custom_method = "writer-custom" reader_actions = _model_actions_mock("reader") reader_actions.some_custom_method = "reader-custom" routed = _RoutedActions(writer_actions, reader_actions, lambda: True) # Unknown method → defaults to writer (safe fallback for write-like ops). assert routed.some_custom_method == "writer-custom" def test_routed_actions_respects_should_use_reader_flag(): """When the routing wrapper marks the reader unavailable, _RoutedActions must redirect reads to the writer instead — without needing to re-fetch the actions accessor.""" from litellm.proxy.db.routing_prisma_wrapper import _RoutedActions writer_actions = _model_actions_mock("writer") reader_actions = _model_actions_mock("reader") use_reader = {"value": True} routed = _RoutedActions(writer_actions, reader_actions, lambda: use_reader["value"]) # Reader healthy → reads to reader. assert routed.find_many is reader_actions.find_many # Reader degrades mid-flight → next read goes to writer. use_reader["value"] = False assert routed.find_many is writer_actions.find_many # --------------------------------------------------------------------------- # Reader graceful degradation # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_connect_swallows_reader_failure_and_falls_back_to_writer(): """A reader connect failure must NOT abort proxy startup. The wrapper flips into degraded mode so subsequent reads route to the writer.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.connect = AsyncMock() reader_inner.connect = AsyncMock(side_effect=RuntimeError("reader unreachable")) routing = RoutingPrismaWrapper(writer=writer, reader=reader) # Must not raise — reader failure is non-fatal. await routing.connect() assert routing.reader_unavailable is True writer_inner.connect.assert_awaited_once() reader_inner.connect.assert_awaited_once() @pytest.mark.asyncio async def test_reads_route_to_writer_when_reader_unavailable(): """Top-level read methods and per-model reads must fall through to the writer while the reader is degraded.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.litellm_usertable = _model_actions_mock("writer_users") reader_inner.litellm_usertable = _model_actions_mock("reader_users") routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._reader_unavailable = True # Top-level reads → writer. assert routing.query_raw is writer_inner.query_raw assert routing.query_first is writer_inner.query_first # Per-model reads → writer actions. actions = routing.litellm_usertable assert actions.find_many is writer_inner.litellm_usertable.find_many assert actions.find_unique is writer_inner.litellm_usertable.find_unique @pytest.mark.asyncio async def test_recreate_prisma_client_recreates_both_writer_and_reader(): """Writer reconnect path calls recreate_prisma_client. The routing wrapper must recreate BOTH clients so a DB-wide event doesn't leave a stale reader.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock() reader = MagicMock() reader.iam_token_db_auth = False reader.recreate_prisma_client = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) with patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}): await routing.recreate_prisma_client("writer-url", http_client=None) writer.recreate_prisma_client.assert_awaited_once_with( "writer-url", http_client=None, expected_generation=None ) reader.recreate_prisma_client.assert_awaited_once_with( "reader-url", http_client=None ) assert routing.reader_unavailable is False @pytest.mark.asyncio async def test_recreate_recovers_reader_after_prior_degradation(): """If a previous connect/recreate degraded the reader, a successful recreate must clear the flag so reads start hitting the reader again.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock() reader = MagicMock() reader.iam_token_db_auth = False reader.recreate_prisma_client = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._reader_unavailable = True with patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}): await routing.recreate_prisma_client("writer-url") assert routing.reader_unavailable is False @pytest.mark.asyncio async def test_recreate_degrades_reader_if_reader_recreate_fails(): """If the reader recreate fails, writer recreate still succeeds and the routing wrapper degrades (does not raise).""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock() reader = MagicMock() reader.iam_token_db_auth = False reader.recreate_prisma_client = AsyncMock( side_effect=RuntimeError("reader still down") ) routing = RoutingPrismaWrapper(writer=writer, reader=reader) with patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}): # Must not raise — writer was recreated, reader is best-effort. await routing.recreate_prisma_client("writer-url") writer.recreate_prisma_client.assert_awaited_once() assert routing.reader_unavailable is True @pytest.mark.asyncio async def test_recreate_degrades_reader_when_replica_url_missing(): """Non-IAM reader needs DATABASE_URL_READ_REPLICA. If it's missing (configuration drift), the wrapper degrades instead of raising.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock() reader = MagicMock() reader.iam_token_db_auth = False reader.recreate_prisma_client = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) # Ensure env var is absent. with patch.dict(os.environ, {}, clear=False): os.environ.pop("DATABASE_URL_READ_REPLICA", None) await routing.recreate_prisma_client("writer-url") writer.recreate_prisma_client.assert_awaited_once() reader.recreate_prisma_client.assert_not_awaited() assert routing.reader_unavailable is True @pytest.mark.asyncio async def test_recreate_iam_reader_refreshes_token(): """IAM-enabled readers must refresh their token (reader has its own parsed endpoint) and pass the fresh URL to recreate_prisma_client.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock() reader = MagicMock() reader.iam_token_db_auth = True reader.get_rds_iam_token = MagicMock(return_value="postgresql://u:fresh@h:5432/db") reader.recreate_prisma_client = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) await routing.recreate_prisma_client("writer-url") reader.get_rds_iam_token.assert_called_once() reader.recreate_prisma_client.assert_awaited_once_with( "postgresql://u:fresh@h:5432/db", http_client=None ) assert routing.reader_unavailable is False @pytest.mark.asyncio async def test_recreate_degrades_when_iam_token_generation_returns_none(): """If `get_rds_iam_token` returns None (e.g. AWS-side failure), the wrapper must degrade rather than crash — this exercises the explicit `raise RuntimeError` inside `_recreate_reader`'s IAM branch.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock() reader = MagicMock() reader.iam_token_db_auth = True reader.get_rds_iam_token = MagicMock(return_value=None) reader.recreate_prisma_client = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) await routing.recreate_prisma_client("writer-url") writer.recreate_prisma_client.assert_awaited_once() reader.recreate_prisma_client.assert_not_awaited() assert routing.reader_unavailable is True def test_writer_and_reader_properties_expose_underlying_wrappers(): """The `writer` and `reader` properties are used by PrismaClient.writer_db to smoke-test the writer specifically during reconnect — they must return the exact wrappers passed in.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, _, reader, _ = _make_wrappers() routing = RoutingPrismaWrapper(writer=writer, reader=reader) assert routing.writer is writer assert routing.reader is reader def test_per_model_accessor_falls_back_when_reader_lacks_attr(): """If the reader Prisma client somehow lacks a model accessor that the writer has (older client / partial mock), the wrapper must fall back to the writer accessor instead of raising AttributeError to the caller.""" from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper # Plain class with only the accessor set on the writer side. Using a real # class instead of MagicMock so attribute access raises AttributeError # naturally instead of auto-creating mock attributes. class _PartialPrisma: pass writer_inner = _PartialPrisma() writer_inner.litellm_usertable = _model_actions_mock("writer_users") reader_inner = _PartialPrisma() # deliberately missing litellm_usertable writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False) reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False) routing = RoutingPrismaWrapper(writer=writer, reader=reader) actions = routing.litellm_usertable # Falls back to the writer's accessor verbatim — not a _RoutedActions wrapper. assert actions is writer_inner.litellm_usertable @pytest.mark.asyncio async def test_writer_recreate_passes_http_client_through(monkeypatch): """When PrismaClient is constructed with an http_client, recreate must forward it to the new Prisma() so connection settings persist across reconnects.""" from litellm.proxy.db.prisma_client import PrismaWrapper captured_kwargs: Dict[str, Any] = {} class FakePrisma: def __init__(self, **kwargs): captured_kwargs.update(kwargs) async def connect(self): return None fake_module = MagicMock() fake_module.Prisma = FakePrisma monkeypatch.setitem(sys.modules, "prisma", fake_module) writer = PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=False) sentinel_http = object() await writer.recreate_prisma_client( "postgresql://u:p@h:5432/db", http_client=sentinel_http ) assert captured_kwargs == {"http": sentinel_http} # --------------------------------------------------------------------------- # IAM endpoint parsing + reader IAM refresh # --------------------------------------------------------------------------- def test_parse_iam_endpoint_from_url_extracts_all_fields(): from litellm.proxy.db.prisma_client import parse_iam_endpoint_from_url ep = parse_iam_endpoint_from_url( "postgresql://litellm_user:initial-token@aurora-reader.example.com:6543/litellm?schema=public" ) assert ep.host == "aurora-reader.example.com" assert ep.port == "6543" assert ep.user == "litellm_user" assert ep.name == "litellm" assert ep.schema == "public" def test_parse_iam_endpoint_defaults_port_to_5432_and_skips_schema(): from litellm.proxy.db.prisma_client import parse_iam_endpoint_from_url ep = parse_iam_endpoint_from_url("postgresql://u@host/dbname") assert ep.host == "host" assert ep.port == "5432" assert ep.user == "u" assert ep.name == "dbname" assert ep.schema is None def test_parse_iam_endpoint_rejects_url_without_user_or_dbname(): from litellm.proxy.db.prisma_client import parse_iam_endpoint_from_url with pytest.raises(ValueError, match="missing host or username"): parse_iam_endpoint_from_url("postgresql://host:5432/db") with pytest.raises(ValueError, match="missing database name"): parse_iam_endpoint_from_url("postgresql://u@host:5432/") def test_iam_endpoint_build_url_inserts_token_verbatim(): from litellm.proxy.db.prisma_client import IAMEndpoint # `generate_iam_auth_token` already URL-encodes the presigned token, so # `build_url` must NOT encode again — double-encoding turned `%3D` into # `%253D` and broke RDS auth on the reader path. ep = IAMEndpoint(host="h", port="5432", user="u", name="db", schema="public") pre_encoded_token = "token%2Fwith%3Fweird%26chars%3Dyes" url = ep.build_url(pre_encoded_token) assert url == f"postgresql://u:{pre_encoded_token}@h:5432/db?schema=public" # Sanity check: no `%25` (the encoding of `%`), confirming we didn't re-encode. assert "%25" not in url @pytest.mark.asyncio async def test_iam_refresh_logs_carry_log_prefix(caplog): """When `log_prefix` is set on a PrismaWrapper, every IAM-related log line emitted by that wrapper must start with the prefix so writer and reader can be told apart in interleaved output.""" from litellm.proxy.db.prisma_client import PrismaWrapper wrapper = PrismaWrapper( original_prisma=MagicMock(), iam_token_db_auth=True, log_prefix="[reader]", ) with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): await wrapper.start_token_refresh_task() # Loop emits "RDS IAM token refresh loop started..." on first tick. # Cancel immediately so the loop body runs once and we can assert. await wrapper.stop_token_refresh_task() messages = [r.getMessage() for r in caplog.records] # Both start and stop notifications carry the prefix. assert any( m.startswith("[reader] Started RDS IAM token proactive refresh") for m in messages ) assert any( m.startswith("[reader] Stopped RDS IAM token refresh background task") for m in messages ) def test_get_rds_iam_token_returns_none_when_iam_disabled(): """`get_rds_iam_token` short-circuits to None when iam_token_db_auth is False — covers the early-return guard at the top of the method.""" from litellm.proxy.db.prisma_client import PrismaWrapper wrapper = PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=False) assert wrapper.get_rds_iam_token() is None @pytest.mark.asyncio async def test_getattr_does_not_block_inside_running_loop_on_expired_token(monkeypatch): """When `__getattr__` runs inside a running event loop and the IAM token is expired, it MUST schedule the refresh as a background task and return immediately. The previous `run_coroutine_threadsafe` + `future.result()` pattern deadlocks the loop (loop thread blocks waiting for a coroutine that needs the loop to run) and times out at 30s — exactly what was breaking the reader on first query.""" from litellm.proxy.db.prisma_client import PrismaWrapper # Stale URL — `is_token_expired` returns True because the password isn't # a parseable IAM token, so we exercise the expired branch. monkeypatch.setenv( "DATABASE_URL_READ_REPLICA", "postgresql://reader:placeholder@reader.aurora.local:5432/litellm", ) inner = MagicMock() inner.query_raw = MagicMock(name="query_raw_attr") wrapper = PrismaWrapper( original_prisma=inner, iam_token_db_auth=True, db_url_env_var="DATABASE_URL_READ_REPLICA", ) # Replace the heavy refresh coroutine with a no-op AsyncMock so we can # observe whether it was scheduled without actually doing the recreate. refresh_calls = {"count": 0} async def fake_refresh(): refresh_calls["count"] += 1 monkeypatch.setattr(wrapper, "_safe_refresh_token", fake_refresh) # Direct attribute access from inside this async test runs __getattr__ # on the loop thread, exercising the in-loop branch. If the previous # `run_coroutine_threadsafe` + `future.result()` pattern were back, this # line would deadlock the loop and the test would hang (and pytest's # per-test timeout would catch it). attr = wrapper.query_raw # Yield once so the scheduled refresh task gets a chance to run. await asyncio.sleep(0) assert attr is inner.query_raw assert refresh_calls["count"] == 1 def test_writer_get_rds_iam_token_defaults_port_when_unset(monkeypatch, unset_database_url): """When DATABASE_PORT is unset, the writer must default to the Postgres standard port instead of passing `None` through. Passing None to `generate_iam_auth_token` makes botocore embed the literal string \"None\" in the presigned URL during signing and crashes with `ValueError: Port could not be cast to integer value as 'None'`.""" from litellm.proxy.db.prisma_client import PrismaWrapper monkeypatch.setenv("DATABASE_HOST", "writer.aurora.local") monkeypatch.delenv("DATABASE_PORT", raising=False) monkeypatch.setenv("DATABASE_USER", "litellm") monkeypatch.setenv("DATABASE_NAME", "litellm") monkeypatch.delenv("DATABASE_SCHEMA", raising=False) captured: Dict[str, Any] = {} def fake_generate(db_host=None, db_port=None, db_user=None, *, region=None): captured["port"] = db_port return "TOKEN" fake_module = MagicMock() fake_module.generate_iam_auth_token = fake_generate monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_module) writer = PrismaWrapper( original_prisma=MagicMock(), iam_token_db_auth=True, ) new_url = writer.get_rds_iam_token() assert captured["port"] == "5432" # default applied, NOT None assert ":5432/litellm" in (new_url or "") def test_writer_get_rds_iam_token_uses_database_host_env_vars(monkeypatch, unset_database_url): """Writer's IAM path (no iam_endpoint configured) reads host/port/user/db from the legacy DATABASE_HOST/PORT/USER/NAME env vars and writes the URL back to DATABASE_URL — this is the pre-read-replica behavior the patch must preserve.""" from litellm.proxy.db.prisma_client import PrismaWrapper monkeypatch.setenv("DATABASE_HOST", "writer.aurora.local") monkeypatch.setenv("DATABASE_PORT", "5432") monkeypatch.setenv("DATABASE_USER", "litellm") monkeypatch.setenv("DATABASE_NAME", "litellm") monkeypatch.setenv("DATABASE_SCHEMA", "public") captured: Dict[str, Any] = {} def fake_generate(db_host=None, db_port=None, db_user=None, *, region=None): captured["host"] = db_host captured["port"] = db_port captured["user"] = db_user return "WRITER-TOKEN" fake_module = MagicMock() fake_module.generate_iam_auth_token = fake_generate monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_module) writer = PrismaWrapper( original_prisma=MagicMock(), iam_token_db_auth=True, # No iam_endpoint → legacy DATABASE_HOST/etc. path. ) new_url = writer.get_rds_iam_token() assert captured == { "host": "writer.aurora.local", "port": "5432", "user": "litellm", } assert new_url == ( "postgresql://litellm:WRITER-TOKEN@writer.aurora.local:5432/litellm?schema=public" ) # Writer updates its own env var (DATABASE_URL by default), not the reader's. assert os.environ["DATABASE_URL"] == new_url def test_reader_iam_refresh_uses_parsed_endpoint(monkeypatch): """The reader generates fresh tokens against its parsed endpoint and writes the new URL to DATABASE_URL_READ_REPLICA — not DATABASE_URL.""" from litellm.proxy.db.prisma_client import IAMEndpoint, PrismaWrapper # Pre-seed env vars so we can prove the reader does NOT touch DATABASE_URL. monkeypatch.setenv("DATABASE_URL", "writer-url-untouched") monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "stale-reader-url") captured: Dict[str, Any] = {} def fake_generate(db_host=None, db_port=None, db_user=None, *, region=None): captured["host"] = db_host captured["port"] = db_port captured["user"] = db_user return "FRESH-TOKEN" fake_module = MagicMock() fake_module.generate_iam_auth_token = fake_generate monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_module) endpoint = IAMEndpoint( host="reader.aurora.local", port="5432", user="lit", name="litellm", schema=None, ) reader = PrismaWrapper( original_prisma=MagicMock(), iam_token_db_auth=True, db_url_env_var="DATABASE_URL_READ_REPLICA", iam_endpoint=endpoint, recreate_uses_datasource=True, ) new_url = reader.get_rds_iam_token() # IAM token generator was called with the reader's parsed endpoint, not # the writer's DATABASE_HOST/PORT/USER env vars. assert captured == { "host": "reader.aurora.local", "port": "5432", "user": "lit", } assert new_url is not None assert new_url.startswith( "postgresql://lit:FRESH-TOKEN@reader.aurora.local:5432/litellm" ) # The reader updates its OWN env var; writer's DATABASE_URL is left alone. assert os.environ["DATABASE_URL_READ_REPLICA"] == new_url assert os.environ["DATABASE_URL"] == "writer-url-untouched" @pytest.mark.asyncio async def test_reader_recreate_uses_datasource_override(monkeypatch): """Reader recreate must pass `datasource={"url": ...}` to Prisma() — Prisma only auto-reads DATABASE_URL, so without the override the new reader URL would be silently ignored.""" from litellm.proxy.db.prisma_client import IAMEndpoint, PrismaWrapper captured_kwargs: Dict[str, Any] = {} class FakePrisma: def __init__(self, **kwargs): captured_kwargs.update(kwargs) async def connect(self): return None fake_module = MagicMock() fake_module.Prisma = FakePrisma monkeypatch.setitem(sys.modules, "prisma", fake_module) reader = PrismaWrapper( original_prisma=MagicMock(), iam_token_db_auth=True, db_url_env_var="DATABASE_URL_READ_REPLICA", iam_endpoint=IAMEndpoint(host="h", port="5432", user="u", name="db"), recreate_uses_datasource=True, ) await reader.recreate_prisma_client( "postgresql://u:newtoken@h:5432/db", http_client=None ) assert captured_kwargs == { "datasource": {"url": "postgresql://u:newtoken@h:5432/db"} } @pytest.mark.asyncio async def test_writer_recreate_does_not_use_datasource(monkeypatch): """Writer keeps relying on Prisma reading DATABASE_URL from env — datasource override must NOT leak into the writer path (would override the freshly rotated env var).""" from litellm.proxy.db.prisma_client import PrismaWrapper captured_kwargs: Dict[str, Any] = {} class FakePrisma: def __init__(self, **kwargs): captured_kwargs.update(kwargs) async def connect(self): return None fake_module = MagicMock() fake_module.Prisma = FakePrisma monkeypatch.setitem(sys.modules, "prisma", fake_module) writer = PrismaWrapper( original_prisma=MagicMock(), iam_token_db_auth=True, ) await writer.recreate_prisma_client( "postgresql://u:newtoken@h:5432/db", http_client=None ) assert "datasource" not in captured_kwargs def test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails( monkeypatch, caplog ): """A transient AWS STS error (or any other failure) during the reader IAM token mint must NOT abort proxy startup. The reader is opt-in, so `PrismaClient.__init__` should log a warning and fall back to the writer-only `PrismaWrapper`. The runtime contract in `RoutingPrismaWrapper.connect` already says reader-side failures are non-fatal — but that code never runs if construction throws first.""" from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") monkeypatch.setenv( "DATABASE_URL_READ_REPLICA", "postgresql://reader_user@reader.aurora.local:5432/litellm", ) class FakePrisma: def __init__(self, **kwargs): self.kwargs = kwargs async def connect(self): return None fake_prisma_module = MagicMock() fake_prisma_module.Prisma = FakePrisma monkeypatch.setitem(sys.modules, "prisma", fake_prisma_module) fake_iam_module = MagicMock() def boom(**_kwargs): raise RuntimeError("simulated AWS STS hiccup") fake_iam_module.generate_iam_auth_token = boom monkeypatch.setitem( sys.modules, "litellm.proxy.auth.rds_iam_token", fake_iam_module ) from litellm.proxy.utils import PrismaClient with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): client = PrismaClient( database_url="postgresql://writer@writer.aurora.local:5432/litellm", proxy_logging_obj=MagicMock(), ) # Construction did not raise, and the proxy is in writer-only mode — # NOT a RoutingPrismaWrapper, so reads will go to the writer. assert isinstance(client.db, PrismaWrapper) assert not isinstance(client.db, RoutingPrismaWrapper) # And the operator gets a clear warning. assert any( "Failed to initialize read replica Prisma client" in r.getMessage() for r in caplog.records ) def test_prisma_client_init_keeps_reader_tls_params_on_the_minted_iam_url( monkeypatch: pytest.MonkeyPatch, ): """The initial reader mint rebuilds the URL from host/port/user/db, so the Prisma TLS dialect on DATABASE_URL_READ_REPLICA must be carried over or a verify-only database rejects the reader and reads fall to the writer.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") monkeypatch.setenv( "DATABASE_URL_READ_REPLICA", "postgresql://reader_user@reader.aurora.local:5432/litellm" "?schema=tenant&sslmode=require&sslcert=/certs/root.pem&sslaccept=strict", ) prisma_factory: Final = MagicMock(name="Prisma") fake_prisma_module: Final = MagicMock(Prisma=prisma_factory) monkeypatch.setitem(sys.modules, "prisma", fake_prisma_module) fake_iam_module: Final = MagicMock(generate_iam_auth_token=MagicMock(return_value="READER-TOKEN")) monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_iam_module) from litellm.proxy.utils import PrismaClient client: Final = PrismaClient( database_url="postgresql://writer@writer.aurora.local:5432/litellm", proxy_logging_obj=MagicMock(), ) assert isinstance(client.db, RoutingPrismaWrapper) reader_url: Final = os.environ["DATABASE_URL_READ_REPLICA"] assert reader_url.startswith("postgresql://reader_user:READER-TOKEN@reader.aurora.local:5432/litellm?") assert parse_qs(urlsplit(reader_url).query) == { "schema": ["tenant"], "sslmode": ["require"], "sslcert": ["/certs/root.pem"], "sslaccept": ["strict"], } assert prisma_factory.call_args_list == [call(), call(datasource={"url": reader_url})] @pytest.mark.asyncio async def test_connect_degrades_writer_when_reader_available(): """A writer connect failure with a healthy reader must NOT abort proxy startup (LIT-3792): startup swallows the raise when allow_requests_on_db_unavailable is set, leaving the proxy with no Prisma client at all, so DB-stored models never load and every request 400s. Degrading instead keeps reads (key auth, model loads) on the replica.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.connect = AsyncMock(side_effect=RuntimeError("primary unreachable")) reader_inner.connect = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) # Must not raise — writer failure is non-fatal while the reader is up. await routing.connect() assert routing.writer_unavailable is True assert routing.reader_unavailable is False writer_inner.connect.assert_awaited_once() reader_inner.connect.assert_awaited_once() # Reads keep routing to the reader. assert routing.query_raw is reader_inner.query_raw @pytest.mark.asyncio async def test_connect_raises_when_writer_and_reader_both_fail(): """Full DB outage: with neither side reachable the wrapper must raise the writer's error so existing allow_requests_on_db_unavailable startup handling applies unchanged.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.connect = AsyncMock(side_effect=RuntimeError("primary down")) reader_inner.connect = AsyncMock(side_effect=RuntimeError("replica down")) routing = RoutingPrismaWrapper(writer=writer, reader=reader) with pytest.raises(RuntimeError, match="primary down"): await routing.connect() @pytest.mark.asyncio async def test_connect_logs_writer_degradation(caplog): """Operators need a clear signal that the proxy booted without a writer.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer, writer_inner, reader, reader_inner = _make_wrappers() writer_inner.connect = AsyncMock(side_effect=RuntimeError("primary unreachable")) reader_inner.connect = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): await routing.connect() assert any( "Failed to connect to primary (writer) DB" in r.getMessage() for r in caplog.records ) @pytest.mark.asyncio async def test_recreate_clears_writer_unavailable(): """A successful writer recreate (health watchdog reconnect once the primary is back) must clear the degraded-writer flag.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock(return_value=True) reader = MagicMock() reader.iam_token_db_auth = False reader.recreate_prisma_client = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._writer_unavailable = True with patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}): await routing.recreate_prisma_client("writer-url") assert routing.writer_unavailable is False @pytest.mark.asyncio async def test_recreate_keeps_writer_unavailable_when_writer_recreate_fails(): """While the primary is still down, a failed writer recreate must leave the degraded flag set so the watchdog keeps retrying.""" from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() writer.recreate_prisma_client = AsyncMock( side_effect=RuntimeError("primary still down") ) reader = MagicMock() reader.iam_token_db_auth = False reader.recreate_prisma_client = AsyncMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._writer_unavailable = True with ( patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}), pytest.raises(RuntimeError, match="primary still down"), ): await routing.recreate_prisma_client("writer-url") assert routing.writer_unavailable is True def test_prisma_client_premints_an_entra_token_for_the_reader(monkeypatch): """Under Azure Entra auth the reader has to be pre-minted the same way the RDS reader already is: Prisma is constructed with a `datasource` URL, so a reader built from the operator's placeholder URL would never carry a real token.""" from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.db.token_auth import AzureEntraTokenAuth monkeypatch.setenv("AZURE_POSTGRESQL_AUTH", "true") monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False) monkeypatch.setenv( "DATABASE_URL_READ_REPLICA", "postgresql://litellm%40contoso.com@reader.postgres.database.azure.com:5432/litellm", ) captured_kwargs: Dict[str, Any] = {} class FakePrisma: def __init__(self, **kwargs): captured_kwargs.update(kwargs) async def connect(self): return None fake_prisma_module = MagicMock() fake_prisma_module.Prisma = FakePrisma monkeypatch.setitem(sys.modules, "prisma", fake_prisma_module) with patch( "litellm.secret_managers.get_azure_ad_token_provider.get_azure_ad_token_provider", return_value=lambda: "ENTRA-TOKEN", ): from litellm.proxy.utils import PrismaClient client = PrismaClient( database_url="postgresql://litellm@writer.postgres.database.azure.com:5432/litellm", proxy_logging_obj=MagicMock(), ) assert isinstance(client.db, RoutingPrismaWrapper) assert captured_kwargs["datasource"] == { "url": "postgresql://litellm%40contoso.com:ENTRA-TOKEN@reader.postgres.database.azure.com:5432/litellm" } assert os.environ["DATABASE_URL_READ_REPLICA"] == captured_kwargs["datasource"]["url"] assert isinstance(client.db._reader.token_auth, AzureEntraTokenAuth) assert isinstance(client.db._writer.token_auth, AzureEntraTokenAuth) assert isinstance(client.db._writer, PrismaWrapper) @pytest.mark.parametrize("opaque_hosts", [False, True]) @pytest.mark.parametrize( ("writer_override", "reader_override"), [(None, None), ("eu-west-1", None), (None, "us-west-2"), (" eu-west-1 ", "us-west-2"), ("", " ")], ) def test_rds_regions_survive_reader_url_startup_and_refresh( monkeypatch: pytest.MonkeyPatch, opaque_hosts: bool, writer_override: str | None, reader_override: str | None, ) -> None: from urllib.parse import unquote from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.utils import PrismaClient for key in tuple(os.environ): if key.startswith(("AWS_", "DATABASE_", "AZURE_POSTGRESQL_", "LITELLM_PGBOUNCER")): monkeypatch.delenv(key) monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIDEXAMPLE") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") monkeypatch.setenv("AWS_REGION", "ap-south-1") monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") for key, value in (("AWS_RDS_REGION", writer_override), ("AWS_RDS_READ_REPLICA_REGION", reader_override)): if value is not None: monkeypatch.setenv(key, value) writer_host: Final = "writer.internal" if opaque_hosts else "w.abc.us-east-1.rds.amazonaws.com" reader_host: Final = "reader.internal" if opaque_hosts else "r.abc.ap-northeast-1.rds.amazonaws.com" monkeypatch.setenv("DATABASE_HOST", writer_host) monkeypatch.setenv("DATABASE_USER", "iam_user") monkeypatch.setenv("DATABASE_NAME", "litellm") monkeypatch.setenv("DATABASE_URL", "") monkeypatch.setenv("DATABASE_URL_READ_REPLICA", f"postgresql://iam_user@{reader_host}:5432/litellm?sslmode=require") monkeypatch.setitem(sys.modules, "prisma", MagicMock()) writer_url: Final = DatabaseURLSettings.from_env().build_writer_url() assert writer_url is not None monkeypatch.setenv("DATABASE_URL", writer_url) client: Final = PrismaClient(database_url=writer_url, proxy_logging_obj=MagicMock()) assert isinstance(client.db, RoutingPrismaWrapper) writer_region: Final = (writer_override or "").strip() or ("ap-south-1" if opaque_hosts else "us-east-1") reader_region: Final = (reader_override or "").strip() or ("ap-south-1" if opaque_hosts else "ap-northeast-1") for url, host, region in ( (writer_url, writer_host, writer_region), (os.environ["DATABASE_URL_READ_REPLICA"], reader_host, reader_region), (client.db.writer.get_rds_iam_token(), writer_host, writer_region), (client.db.reader.get_rds_iam_token(), reader_host, reader_region), ): assert url is not None password: Final = urlsplit(url).password assert password is not None token: Final = unquote(password) assert token.split("?", 1)[0] == f"{host}:5432/" assert parse_qs(urlsplit(token).query)["X-Amz-Credential"][0].split("/")[2] == region assert parse_qs(urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query)["sslmode"] == ["require"] assert os.environ["AWS_REGION"] == "ap-south-1"