mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(proxy): read the user row past the recent-miss memo on a database-only lookup
get_user_object skipped the database for db_cache_expiry seconds after a miss on the same worker even when the caller asked for check_db_only, so the token exchange mint could answer no_active_key for a user JWT auth had just created. A database-only read now always reaches the database.
This commit is contained in:
parent
788e6eb737
commit
f6ee046199
2 changed files with 28 additions and 1 deletions
|
|
@ -2519,7 +2519,7 @@ async def get_user_object(
|
|||
raise Exception("No db connected")
|
||||
try:
|
||||
db_access_time_key: Final = f"user_id:{user_id}"
|
||||
should_check_db: Final = _should_check_db(
|
||||
should_check_db: Final = bool(check_db_only) or _should_check_db(
|
||||
key=db_access_time_key,
|
||||
last_db_access_time=last_db_access_time,
|
||||
db_cache_expiry=db_cache_expiry,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -915,6 +916,32 @@ async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context(
|
|||
assert isinstance(exc_info.value.__context__, ConnectionError)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_user_object_check_db_only_ignores_recent_miss(monkeypatch):
|
||||
"""A database-only read is never answered by the per-worker negative memo: a row created after a miss on
|
||||
this worker is returned within db_cache_expiry seconds instead of raising UserNotFoundError, so the token
|
||||
exchange mints for a user JWT auth just accepted."""
|
||||
from litellm.proxy.auth import auth_checks
|
||||
|
||||
user_id = "memo-probe-user"
|
||||
monkeypatch.setitem(auth_checks.last_db_access_time, f"user_id:{user_id}", (None, time.time()))
|
||||
db_row = LiteLLM_UserTable(user_id=user_id, user_email=None, user_role="internal_user")
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=db_row)
|
||||
|
||||
result = await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
user_id_upsert=False,
|
||||
check_db_only=True,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.user_id == user_id
|
||||
mock_prisma_client.db.litellm_usertable.find_unique.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_user_object_upsert_includes_user_email():
|
||||
"""Test that user_email is included when creating a new user via get_user_object upsert"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue