refactor(db): route the spend capture rate read through replica_db

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/spend_tracking/spend_capture_rate.py

@@ -86,7 +86,7 @@ async def captured_spend_by_day(
     end_date: date,
 ) -> Mapping[str, float]:
     """LiteLLM's tracked spend per UTC day (ISO date) for the given ``custom_llm_provider`` values."""
-    rows: Final = await prisma_client.db.query_raw(
+    rows: Final = await prisma_client.replica_db.query_raw(
         _CAPTURED_SPEND_BY_DAY_SQL, start_date.isoformat(), end_date.isoformat(), tuple(litellm_providers)
     )
     return MappingProxyType({row.date: row.spend for row in _CAPTURED_SPEND_ROWS.validate_python(rows)})

tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py

@@ -733,6 +733,7 @@ async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted():
     with the redaction marker."""
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
 
     mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
         return_value=SimpleNamespace(
@@ -830,6 +831,7 @@ async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_
     redaction marker, not just top-level secret keys."""
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
 
     mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
         return_value=SimpleNamespace(
@@ -967,6 +969,7 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted():
     the stored credential."""
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
 
     mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
         return_value={
@@ -1165,6 +1168,7 @@ _KILL_SWITCH: Final = {
 async def test_add_agent_to_db_stores_kill_switch_json_and_a_json_null_when_unset():
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
     mock_create = AsyncMock(return_value=_agent_row_mock([]))
     mock_prisma.db.litellm_agentstable.create = mock_create
 
@@ -1191,6 +1195,7 @@ async def test_add_agent_to_db_stores_kill_switch_json_and_a_json_null_when_unse
 async def test_add_agent_to_db_rejects_a_kill_switch_with_a_non_http_url():
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
     mock_prisma.db.litellm_agentstable.create = AsyncMock(return_value=_agent_row_mock([]))
 
     with pytest.raises(Exception, match="absolute http"):
@@ -1210,6 +1215,7 @@ async def test_add_agent_to_db_rejects_a_kill_switch_with_a_non_http_url():
 async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on_null():
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
     mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
         return_value={
             "agent_id": "agent-123",
@@ -1239,6 +1245,7 @@ async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on
 async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_the_marker():
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
     mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
         return_value={
             "agent_id": "agent-123",
@@ -1274,6 +1281,7 @@ async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_t
 async def test_update_agent_in_db_clears_kill_switch_when_omitted_and_restores_secret_when_echoed():
     registry: Final = AgentRegistry()
     mock_prisma: Final = MagicMock()
+    mock_prisma.replica_db = mock_prisma.db
     mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
         return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=json.dumps(_KILL_SWITCH))
     )

tests/test_litellm/proxy/agent_endpoints/test_endpoints.py

@@ -1330,6 +1330,7 @@ def test_kill_switch_trigger_fires_the_db_row_config_over_a_stale_in_memory_copy
         kill_switch={"url": "https://ops.example.com/kill-v2", "method": "DELETE", "auth": None},
     )
     prisma: Final = MagicMock()
+    prisma.replica_db = prisma.db
     prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=db_row)
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
     fake: Final = _FakeKillSwitchClient(httpx.Response(204))
@@ -1353,6 +1354,7 @@ def test_get_agent_redacts_kill_switch_secret_for_admins_and_hides_it_from_other
 
     def _get_as(role: LitellmUserRoles):
         with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
+            mock_prisma.replica_db = mock_prisma.db
             mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
             mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
             return _make_app_with_role(role).get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"})

tests/test_litellm/proxy/spend_tracking/test_spend_capture_rate.py

@@ -68,6 +68,7 @@ def _page(*buckets: dict[str, object], next_page: str | None = None) -> dict[str
 
 def _fake_prisma(rows: list[dict[str, object]]) -> MagicMock:
     prisma = MagicMock()
+    prisma.replica_db = prisma.db
     prisma.db.query_raw = AsyncMock(return_value=rows)
     return prisma
 
@@ -306,6 +307,7 @@ async def test_a_healthy_scheduled_check_publishes_and_never_touches_the_alert_l
 async def test_a_scheduled_check_that_fails_never_claims_the_alert_window(monkeypatch):
     lock = _pod_lock(acquired=True)
     prisma = MagicMock()
+    prisma.replica_db = prisma.db
     prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("database gone"))
 
     with pytest.raises(RuntimeError, match="database gone"):
@@ -395,6 +397,7 @@ class _PsycopgPrisma:
 
     def __init__(self, conn: psycopg.Connection) -> None:
         self.db = self
+        self.replica_db = self
         self._conn = conn
 
     async def query_raw(self, sql: str, *params: object) -> list[dict[str, object]]: