mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-08 22:25:25 +00:00
fix(budget_reset): use raw SQL for IS NOT NULL filter on Json? columns
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`.
This commit is contained in:
Yuneng Jiang
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user