From c41567eaa0caaffc6070b2806adc6369b61fe017 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 23 Apr 2026 12:26:25 -0700 Subject: [PATCH] fix(budget_reset): use raw SQL for IS NOT NULL filter on Json? columns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The periodic budget-window reset job filtered keys/teams with `where={"budget_limits": {"not": None}}`. The prisma-client-python library does not support null-filtering on `Json?` columns (no DbNull/JsonNull sentinel — upstream issue #714). The client drops the `None` value during serialization and the engine rejects the query with `MissingRequiredValueError: where.budget_limits.not: A value is required but not set`, so neither the key nor team reset path runs. Switch those two `find_many` calls to `query_raw` with `WHERE budget_limits IS NOT NULL`, selecting only the PK and the `budget_limits` column. Writes still go through the ORM. Add unit tests covering the expired/unexpired paths for keys and teams, string-encoded JSON payloads, empty payloads, error isolation between the two paths, and a regression guard asserting the query still uses `IS NOT NULL`. --- .../proxy/common_utils/reset_budget_job.py | 34 +-- .../common_utils/test_reset_budget_job.py | 236 +++++++++++++++++- 2 files changed, 253 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index b11af04e2b..e486336cec 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -632,20 +632,27 @@ class ResetBudgetJob: now = datetime.utcnow() + # Note on raw SQL: prisma-client-python does not support null-filtering + # on `Json?` columns (no DbNull/JsonNull sentinel — see + # RobertCraigie/prisma-client-py#714). We use `query_raw` with + # `IS NOT NULL` so we don't materialize every key/team row on each + # tick of the reset job. Writes still go through the ORM. + # --- Keys --- try: - all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many( - where={"budget_limits": {"not": None}} # type: ignore[arg-type] + key_rows = await self.prisma_client.db.query_raw( + 'SELECT token, budget_limits FROM "LiteLLM_VerificationToken" ' + "WHERE budget_limits IS NOT NULL" ) - for key in all_keys: - raw = key.budget_limits # type: ignore[attr-defined] + for row in key_rows: + raw = row["budget_limits"] if not raw: continue windows: list = raw if isinstance(raw, list) else json.loads(raw) changed = False for window in windows: counter_key = ( - f"spend:key:{key.token}:window:{window['budget_duration']}" + f"spend:key:{row['token']}:window:{window['budget_duration']}" ) if await ResetBudgetJob._reset_expired_window( window, counter_key, spend_counter_cache, now @@ -653,7 +660,7 @@ class ResetBudgetJob: changed = True if changed: await self.prisma_client.db.litellm_verificationtoken.update( - where={"token": key.token}, + where={"token": row["token"]}, data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) except Exception as e: @@ -663,26 +670,25 @@ class ResetBudgetJob: # --- Teams --- try: - all_teams = await self.prisma_client.db.litellm_teamtable.find_many( - where={"budget_limits": {"not": None}} # type: ignore[arg-type] + team_rows = await self.prisma_client.db.query_raw( + 'SELECT team_id, budget_limits FROM "LiteLLM_TeamTable" ' + "WHERE budget_limits IS NOT NULL" ) - for team in all_teams: - raw = team.budget_limits # type: ignore[attr-defined] + for row in team_rows: + raw = row["budget_limits"] if not raw: continue windows = raw if isinstance(raw, list) else json.loads(raw) changed = False for window in windows: - counter_key = ( - f"spend:team:{team.team_id}:window:{window['budget_duration']}" - ) + counter_key = f"spend:team:{row['team_id']}:window:{window['budget_duration']}" if await ResetBudgetJob._reset_expired_window( window, counter_key, spend_counter_cache, now ): changed = True if changed: await self.prisma_client.db.litellm_teamtable.update( - where={"team_id": team.team_id}, + where={"team_id": row["team_id"]}, data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) except Exception as e: diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 32f043be5b..379ccf4d9a 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1,7 +1,9 @@ import asyncio +import json import os import sys import time +import types from datetime import datetime, timedelta, timezone from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock @@ -696,9 +698,9 @@ def test_reset_budget_resets_endusers_with_null_budget_id( # Both end users should have been reset updated = mock_prisma_client.updated_data["enduser"] - assert len(updated) == 2, ( - f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" - ) + assert ( + len(updated) == 2 + ), f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" user_ids = {u.user_id for u in updated} assert "enduser-explicit" in user_ids @@ -819,3 +821,231 @@ def test_reset_budget_for_team_members_preserves_total_spend(): assert call_kwargs["where"]["budget_id"]["in"] == ["budget-1"] assert call_kwargs["data"] == {"spend": 0} assert "total_spend" not in call_kwargs["data"] + + +# --------------------------------------------------------------------------- +# reset_budget_windows (per-key / per-team concurrent window resets) +# --------------------------------------------------------------------------- + + +def _make_reset_budget_windows_job( + monkeypatch, + key_rows: List[Dict[str, Any]], + team_rows: List[Dict[str, Any]], +): + """Build a ResetBudgetJob with a fully-mocked prisma client and a fake + `litellm.proxy.proxy_server` module exposing a stub `spend_counter_cache`. + + Returns (job, prisma_client_mock, spend_counter_cache_mock). + """ + prisma_client = MagicMock() + + async def fake_query_raw(query: str, *args, **kwargs): + # Dispatch by table name in the SQL so a single stub covers both calls. + if '"LiteLLM_VerificationToken"' in query: + return key_rows + if '"LiteLLM_TeamTable"' in query: + return team_rows + raise AssertionError(f"Unexpected query_raw call: {query}") + + prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) + prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) + + # Stub out litellm.proxy.proxy_server so the in-function + # `from litellm.proxy.proxy_server import spend_counter_cache` resolves + # without importing the real (heavy) module. + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = None # skip the async redis branch + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + return job, prisma_client, spend_counter_cache + + +def test_reset_budget_windows_uses_is_not_null_filter(monkeypatch): + """Regression guard for the Prisma client limitation documented in + RobertCraigie/prisma-client-py#714: `{"not": None}` on a `Json?` column + raises `MissingRequiredValueError`. We work around it by using `query_raw` + with `IS NOT NULL`. If someone reverts to the ORM filter, this test fails. + """ + job, prisma_client, _ = _make_reset_budget_windows_job( + monkeypatch, key_rows=[], team_rows=[] + ) + + asyncio.run(job.reset_budget_windows()) + + queries = [call.args[0] for call in prisma_client.db.query_raw.await_args_list] + assert len(queries) == 2, queries + key_query, team_query = queries + + assert '"LiteLLM_VerificationToken"' in key_query + assert "budget_limits IS NOT NULL" in key_query + assert '"LiteLLM_TeamTable"' in team_query + assert "budget_limits IS NOT NULL" in team_query + + +def test_reset_budget_windows_resets_expired_key_window(monkeypatch): + """A key whose window's `reset_at` has passed gets an update with a new + `reset_at` in the future, and the in-memory spend counter is cleared.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + + asyncio.run(job.reset_budget_windows()) + + # Update should have been called exactly once with the expired token. + prisma_client.db.litellm_verificationtoken.update.assert_awaited_once() + call_kwargs = prisma_client.db.litellm_verificationtoken.update.await_args.kwargs + assert call_kwargs["where"] == {"token": "sk-expired"} + + # The `budget_limits` payload is re-serialized JSON with a bumped reset_at. + written_windows = json.loads(call_kwargs["data"]["budget_limits"]) + assert len(written_windows) == 1 + new_reset_at = datetime.fromisoformat( + written_windows[0]["reset_at"].replace("Z", "+00:00") + ).replace(tzinfo=None) + assert new_reset_at > now + + # The spend counter for this key+window was cleared. + spend_counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:key:sk-expired:window:1d", value=0.0 + ) + + +def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): + """If `reset_at` is in the future, no write should happen for that key.""" + now = datetime.utcnow() + future = (now + timedelta(hours=1)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-future", + "budget_limits": [{"budget_duration": "1d", "reset_at": future}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_not_awaited() + + +def test_reset_budget_windows_resets_expired_team_window(monkeypatch): + """Same as the key test, but for teams.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=1)).isoformat() + "Z" + + team_rows = [ + { + "team_id": "team-expired", + "budget_limits": [{"budget_duration": "30d", "reset_at": expired}], + } + ] + job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job( + monkeypatch, key_rows=[], team_rows=team_rows + ) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_teamtable.update.assert_awaited_once() + call_kwargs = prisma_client.db.litellm_teamtable.update.await_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-expired"} + assert "budget_limits" in call_kwargs["data"] + + spend_counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:team:team-expired:window:30d", value=0.0 + ) + + +def test_reset_budget_windows_handles_string_budget_limits(monkeypatch): + """Defensive: if `query_raw` returns `budget_limits` as a JSON-encoded + string (driver-dependent), the code still parses and resets it. + """ + now = datetime.utcnow() + expired = (now - timedelta(minutes=1)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-string-limits", + "budget_limits": json.dumps( + [{"budget_duration": "1d", "reset_at": expired}] + ), + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_awaited_once() + + +def test_reset_budget_windows_skips_row_with_empty_budget_limits(monkeypatch): + """A row whose `budget_limits` comes back as an empty/falsy payload + (shouldn't happen given the WHERE filter, but we guard anyway) must not + trigger an update or crash the loop.""" + key_rows = [ + {"token": "sk-empty-list", "budget_limits": []}, + {"token": "sk-empty-str", "budget_limits": ""}, + ] + job, prisma_client, _ = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_not_awaited() + + +def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch): + """If the key query raises, the teams path still runs (and vice-versa). + Each side has its own try/except; this locks that in.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=1)).isoformat() + "Z" + + prisma_client = MagicMock() + + async def fake_query_raw(query: str, *args, **kwargs): + if '"LiteLLM_VerificationToken"' in query: + raise RuntimeError("boom") + if '"LiteLLM_TeamTable"' in query: + return [ + { + "team_id": "team-ok", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + raise AssertionError(query) + + prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) + + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = None + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + + asyncio.run(job.reset_budget_windows()) # must not raise + + prisma_client.db.litellm_teamtable.update.assert_awaited_once()