Merge pull request #41283 from BerriAI/litellm_writer_pinned_raw_sql_writes

fix(proxy): keep access-group raw SQL writes on the writer while writer_unavailable is stale

litellm/proxy/db/routing_prisma_wrapper.py

@@ -81,6 +81,11 @@ class WriterPinnedClient:
         self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db
 
 
+def writer_wrapper(db: "PrismaWrapper | RoutingPrismaWrapper") -> PrismaWrapper:
+    """Unlike `WriterPinnedClient`, ignores `writer_unavailable`: a raw SQL write has no replica fallback."""
+    return db.writer if isinstance(db, RoutingPrismaWrapper) else db
+
+
 class RoutingPrismaWrapper:
     """
     Routes Prisma operations between a writer and a reader Prisma client.

litellm/proxy/management_helpers/access_group_key_sync.py

@@ -38,7 +38,7 @@ from litellm.proxy._types import (
 from litellm.proxy.auth.auth_checks import (
     _delete_cache_access_object,  # pyright: ignore[reportPrivateUsage]  # the access-group endpoints reach for this same cache primitive
 )
-from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
+from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
 from litellm.repositories.table_repositories import AccessGroupRepository
 
 
@@ -75,7 +75,7 @@ _REPOINT_KEY_SQL: Final = (
 def _raw_executor(prisma_client: object) -> _RawExecutor:
     """Narrow the untyped Prisma client down to the raw-query call this module makes, pinned to the writer."""
     db: Final = AccessGroupRepository(prisma_client).prisma_client.db  # pyright: ignore[reportAny]  # untyped Prisma client
-    return WriterPinnedClient(db).db  # pyright: ignore[reportAny, reportReturnType]  # untyped Prisma client behind the pin
+    return writer_wrapper(db)  # pyright: ignore[reportAny, reportReturnType]  # untyped Prisma client behind the pin
 
 
 async def _invalidate_access_group_cache(access_group_id: str) -> None:

litellm/proxy/management_helpers/access_group_model_sync.py

@@ -10,7 +10,7 @@ from typing import Final, Protocol
 
 from pydantic import BaseModel
 
-from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
+from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
 from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_caches
 from litellm.repositories.table_repositories import AccessGroupRepository
 from litellm.router import Router
@@ -56,7 +56,7 @@ _REMOVE_MODEL_NAME_SQL: Final = (
 
 def _raw_executor(prisma_client: object) -> _RawExecutor:
     db: Final = AccessGroupRepository(prisma_client).prisma_client.db  # pyright: ignore[reportAny]  # untyped Prisma client
-    return WriterPinnedClient(db).db  # pyright: ignore[reportAny, reportReturnType]  # untyped Prisma client behind the pin
+    return writer_wrapper(db)  # pyright: ignore[reportAny, reportReturnType]  # untyped Prisma client behind the pin
 
 
 def _config_sourced_sibling(llm_router: Router, deployment_id: str, model_id: str) -> bool:

tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py

@@ -145,6 +145,18 @@ def test_writer_pinned_client_yields_to_routed_reads_when_writer_down():
     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

tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py

@@ -11,14 +11,15 @@ from litellm.proxy.management_helpers.access_group_key_sync import (
 )
 
 
-def _routed_prisma_client():
+def _routed_prisma_client(writer_unavailable: bool = False):
     writer_inner = MagicMock(name="writer_prisma")
     reader_inner = MagicMock(name="reader_prisma")
     writer_inner.query_raw = AsyncMock(return_value=[])
-    reader_inner.query_raw = AsyncMock(return_value=[])
+    reader_inner.query_raw = AsyncMock(side_effect=RuntimeError("cannot execute UPDATE in a read-only transaction"))
     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)
+    routing._writer_unavailable = writer_unavailable
     return SimpleNamespace(db=routing), writer_inner, reader_inner
 
 
@@ -39,6 +40,39 @@ async def test_regeneration_repoint_update_runs_on_the_writer():
     reader_inner.query_raw.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_regeneration_repoint_update_stays_on_the_writer_while_writer_flagged_unavailable():
+    prisma_client, writer_inner, reader_inner = _routed_prisma_client(writer_unavailable=True)
+
+    await sync_key_regeneration_access_group_membership(
+        prisma_client=prisma_client,
+        previous_key_token="old-token",
+        new_key_token="new-token",
+        data=None,
+        existing_key_row=MagicMock(),
+    )
+
+    writer_inner.query_raw.assert_awaited_once()
+    assert writer_inner.query_raw.await_args.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"')
+    assert writer_inner.query_raw.await_args.args[1:] == ("old-token", "new-token")
+    reader_inner.query_raw.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_membership_attach_and_detach_updates_stay_on_the_writer_while_writer_flagged_unavailable():
+    prisma_client, writer_inner, reader_inner = _routed_prisma_client(writer_unavailable=True)
+
+    await sync_key_access_group_membership(
+        prisma_client=prisma_client,
+        key_token="token",
+        previous_access_group_ids=["ag-old"],
+        updated_access_group_ids=["ag-new"],
+    )
+
+    assert writer_inner.query_raw.await_count == 2
+    reader_inner.query_raw.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_membership_attach_and_detach_updates_run_on_the_writer():
     prisma_client, writer_inner, reader_inner = _routed_prisma_client()

tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py

@@ -13,7 +13,7 @@ from litellm.proxy.management_helpers.access_group_model_sync import (
 _INVALIDATE = "litellm.proxy.management_helpers.access_group_model_sync.invalidate_access_group_caches"
 
 
-def _routed_prisma_client(deployment_count: int):
+def _routed_prisma_client(deployment_count: int, writer_unavailable: bool = False):
     async def query_raw(sql, *params):
         if sql.startswith("SELECT COUNT(*)"):
             return [{"deployment_count": deployment_count}]
@@ -26,6 +26,7 @@ def _routed_prisma_client(deployment_count: int):
     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)
+    routing._writer_unavailable = writer_unavailable
     return SimpleNamespace(db=routing), writer_inner, reader_inner
 
 
@@ -53,6 +54,20 @@ async def test_rename_replaces_the_old_name_when_no_other_deployment_carries_it(
     reader_inner.query_raw.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_rename_update_stays_on_the_writer_while_writer_flagged_unavailable():
+    prisma_client, writer_inner, reader_inner = _routed_prisma_client(deployment_count=0, writer_unavailable=True)
+
+    with patch(_INVALIDATE, new=AsyncMock()):
+        await sync_access_groups_for_renamed_model(
+            prisma_client, model_id="m-1", old_name="gpt-5.6", new_name="gpt-5.6-eu", llm_router=None
+        )
+
+    (update_call,) = _access_group_updates(writer_inner)
+    assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu")
+    reader_inner.query_raw.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_rename_appends_the_new_name_when_a_sibling_row_keeps_the_old_one():
     prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=1)
@@ -168,3 +183,17 @@ async def test_delete_keeps_the_name_while_a_sibling_row_still_backs_it():
 
     assert _access_group_updates(writer_inner) == []
     invalidate.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_delete_counts_backing_rows_on_the_writer_not_a_lagging_replica_while_writer_flagged_unavailable():
+    prisma_client, writer_inner, reader_inner = _routed_prisma_client(deployment_count=0, writer_unavailable=True)
+    reader_inner.query_raw = AsyncMock(return_value=[{"deployment_count": 1}])
+
+    with patch(_INVALIDATE, new=AsyncMock()) as invalidate:
+        await sync_access_groups_for_deleted_model(prisma_client, model_id="m-1", model_name="gpt-5.6", llm_router=None)
+
+    (update_call,) = _access_group_updates(writer_inner)
+    assert "array_remove" in update_call.args[0]
+    invalidate.assert_awaited_once_with(("ag-1", "ag-2"))
+    reader_inner.query_raw.assert_not_awaited()