mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-09 12:24:21 +00:00
refactor: use DualCache for UI settings reads and get_team_object for team lookup
- Add get_ui_settings_cached() helper that reads from DualCache first, falls back to DB, and populates cache on miss. - Update update_ui_settings() to set cache after DB write so subsequent reads see new values immediately. - Replace raw prisma_client.db.litellm_teamtable.find_unique with the existing get_team_object helper which uses DualCache. - Update all tests to mock get_ui_settings_cached and get_team_object instead of raw DB calls. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
c631708df6
commit
af61132a3f
@@ -30,7 +30,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
get_daily_activity,
|
||||
get_daily_activity_aggregated,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import get_user_object
|
||||
from litellm.proxy.auth.auth_checks import get_team_object, get_user_object
|
||||
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
generate_key_helper_fn,
|
||||
@@ -1869,21 +1869,17 @@ async def ui_view_users(
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
get_ui_settings_cached,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
||||
|
||||
try:
|
||||
# Read the scope_user_search_to_org flag from the DB
|
||||
ui_settings_row = (
|
||||
await prisma_client.db.litellm_uisettings.find_unique(
|
||||
where={"id": "ui_settings"}
|
||||
)
|
||||
)
|
||||
scope_flag = False
|
||||
if ui_settings_row is not None:
|
||||
settings_json = ui_settings_row.settings or {} # type: ignore[union-attr]
|
||||
scope_flag = bool(settings_json.get("scope_user_search_to_org", False))
|
||||
# Read the scope_user_search_to_org flag (cached)
|
||||
ui_settings = await get_ui_settings_cached()
|
||||
scope_flag = bool(ui_settings.get("scope_user_search_to_org", False))
|
||||
|
||||
org_filter_ids: Optional[List[str]] = None
|
||||
|
||||
@@ -1915,27 +1911,29 @@ async def ui_view_users(
|
||||
if org_admin_org_ids:
|
||||
org_filter_ids = org_admin_org_ids
|
||||
elif team_id is not None:
|
||||
# Look up the team to check if it belongs to an org
|
||||
team_row = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
if team_row is not None:
|
||||
team_obj = LiteLLM_TeamTable(**team_row.model_dump())
|
||||
if _is_user_team_admin(user_api_key_dict, team_obj):
|
||||
if team_obj.organization_id:
|
||||
org_filter_ids = [team_obj.organization_id]
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "scope_user_search_to_org is enabled and this team is not part of an organization. Contact your proxy admin to adjust this setting."
|
||||
},
|
||||
)
|
||||
# Look up the team via cached helper
|
||||
try:
|
||||
team_obj = await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except HTTPException:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "scope_user_search_to_org is enabled. Only proxy admins, organization admins, or team admins can search users."
|
||||
},
|
||||
)
|
||||
if _is_user_team_admin(user_api_key_dict, team_obj):
|
||||
if team_obj.organization_id:
|
||||
org_filter_ids = [team_obj.organization_id]
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "scope_user_search_to_org is enabled. Only proxy admins, organization admins, or team admins can search users."
|
||||
"error": "scope_user_search_to_org is enabled and this team is not part of an organization. Contact your proxy admin to adjust this setting."
|
||||
},
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -980,6 +980,48 @@ async def get_in_product_nudges():
|
||||
return InProductNudgeResponse(is_claude_code_enabled=False)
|
||||
|
||||
|
||||
UI_SETTINGS_CACHE_KEY = "ui_settings:settings_dict"
|
||||
|
||||
|
||||
async def get_ui_settings_cached() -> Dict[str, Any]:
|
||||
"""
|
||||
Return the persisted UI settings dict, using DualCache for reads.
|
||||
|
||||
Cache hit → return cached dict immediately.
|
||||
Cache miss → read from DB, populate cache, return dict.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
# 1. Try cache
|
||||
cached = await user_api_key_cache.async_get_cache(key=UI_SETTINGS_CACHE_KEY)
|
||||
if cached is not None and isinstance(cached, dict):
|
||||
return cached
|
||||
|
||||
# 2. Fallback to DB
|
||||
if prisma_client is None:
|
||||
return {}
|
||||
|
||||
db_record = await prisma_client.db.litellm_uisettings.find_unique(
|
||||
where={"id": "ui_settings"}
|
||||
)
|
||||
ui_settings: Dict[str, Any] = {}
|
||||
if db_record and db_record.ui_settings:
|
||||
raw = db_record.ui_settings
|
||||
ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw)
|
||||
|
||||
# Sanitize
|
||||
ui_settings = {
|
||||
k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS
|
||||
}
|
||||
|
||||
# 3. Populate cache
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=UI_SETTINGS_CACHE_KEY, value=ui_settings
|
||||
)
|
||||
|
||||
return ui_settings
|
||||
|
||||
|
||||
@router.get(
|
||||
"/get/ui_settings",
|
||||
tags=["UI Settings"],
|
||||
@@ -1108,6 +1150,16 @@ async def update_ui_settings(
|
||||
|
||||
general_settings.update(_flags_to_sync)
|
||||
|
||||
# Invalidate + set DualCache so subsequent reads see the new values immediately
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
sanitized = {
|
||||
k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS
|
||||
}
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=UI_SETTINGS_CACHE_KEY, value=sanitized
|
||||
)
|
||||
|
||||
return {
|
||||
"message": "UI settings updated successfully",
|
||||
"status": "success",
|
||||
|
||||
@@ -54,11 +54,11 @@ async def test_ui_view_users_with_null_email(mocker, caplog):
|
||||
|
||||
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
|
||||
|
||||
# Flag OFF by default — no settings row
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return None
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
# Flag OFF by default
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={},
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
@@ -92,10 +92,10 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker):
|
||||
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
|
||||
|
||||
# Flag OFF by default
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return None
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={},
|
||||
)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
await ui_view_users(
|
||||
@@ -132,13 +132,10 @@ async def test_ui_view_users_org_admin_filtered_by_org(mocker):
|
||||
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
|
||||
|
||||
# Flag ON
|
||||
mock_settings_row = mocker.MagicMock()
|
||||
mock_settings_row.settings = {"scope_user_search_to_org": True}
|
||||
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return mock_settings_row
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())
|
||||
@@ -185,13 +182,10 @@ async def test_ui_view_users_non_org_admin_returns_403(mocker):
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
|
||||
# Flag ON
|
||||
mock_settings_row = mocker.MagicMock()
|
||||
mock_settings_row.settings = {"scope_user_search_to_org": True}
|
||||
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return mock_settings_row
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())
|
||||
@@ -237,11 +231,11 @@ async def test_ui_view_users_flag_off_internal_user_can_search(mocker):
|
||||
|
||||
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
|
||||
|
||||
# Flag OFF — no settings row
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return None
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
# Flag OFF
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={},
|
||||
)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
response = await ui_view_users(
|
||||
@@ -261,7 +255,7 @@ async def test_ui_view_users_flag_on_team_admin_org_team(mocker):
|
||||
"""
|
||||
Flag ON, team admin for org-bound team: org filter is applied using team's org.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_TeamTable, Member
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
org_id = "org-456"
|
||||
@@ -278,30 +272,26 @@ async def test_ui_view_users_flag_on_team_admin_org_team(mocker):
|
||||
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
|
||||
|
||||
# Flag ON
|
||||
mock_settings_row = mocker.MagicMock()
|
||||
mock_settings_row.settings = {"scope_user_search_to_org": True}
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return mock_settings_row
|
||||
# Mock get_team_object
|
||||
team_obj = LiteLLM_TeamTableCachedObj(
|
||||
team_id=tid,
|
||||
team_alias="test-team",
|
||||
organization_id=org_id,
|
||||
members_with_roles=[{"user_id": "team-admin-user", "role": "admin"}],
|
||||
)
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
async def mock_get_team_object(*args, **kwargs):
|
||||
return team_obj
|
||||
|
||||
# Team lookup
|
||||
mock_team_row = mocker.MagicMock()
|
||||
mock_team_row.model_dump.return_value = {
|
||||
"team_id": tid,
|
||||
"team_alias": "test-team",
|
||||
"organization_id": org_id,
|
||||
"members_with_roles": [{"user_id": "team-admin-user", "role": "admin"}],
|
||||
"admins": [],
|
||||
"members": [],
|
||||
"blocked": False,
|
||||
}
|
||||
|
||||
async def mock_find_unique_team(*args, **kwargs):
|
||||
return mock_team_row
|
||||
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.internal_user_endpoints.get_team_object",
|
||||
side_effect=mock_get_team_object,
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())
|
||||
@@ -337,35 +327,32 @@ async def test_ui_view_users_flag_on_team_admin_non_org_team_403(mocker):
|
||||
Flag ON, team admin for non-org team: returns 403.
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
tid = "team-no-org"
|
||||
|
||||
# Flag ON
|
||||
mock_settings_row = mocker.MagicMock()
|
||||
mock_settings_row.settings = {"scope_user_search_to_org": True}
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return mock_settings_row
|
||||
# Mock get_team_object — team has no organization_id
|
||||
team_obj = LiteLLM_TeamTableCachedObj(
|
||||
team_id=tid,
|
||||
team_alias="no-org-team",
|
||||
organization_id=None,
|
||||
members_with_roles=[{"user_id": "team-admin-user", "role": "admin"}],
|
||||
)
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
async def mock_get_team_object(*args, **kwargs):
|
||||
return team_obj
|
||||
|
||||
# Team lookup — no organization_id
|
||||
mock_team_row = mocker.MagicMock()
|
||||
mock_team_row.model_dump.return_value = {
|
||||
"team_id": tid,
|
||||
"team_alias": "no-org-team",
|
||||
"organization_id": None,
|
||||
"members_with_roles": [{"user_id": "team-admin-user", "role": "admin"}],
|
||||
"admins": [],
|
||||
"members": [],
|
||||
"blocked": False,
|
||||
}
|
||||
|
||||
async def mock_find_unique_team(*args, **kwargs):
|
||||
return mock_team_row
|
||||
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.internal_user_endpoints.get_team_object",
|
||||
side_effect=mock_get_team_object,
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())
|
||||
@@ -409,13 +396,10 @@ async def test_ui_view_users_flag_on_non_admin_no_team_id_403(mocker):
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
|
||||
# Flag ON
|
||||
mock_settings_row = mocker.MagicMock()
|
||||
mock_settings_row.settings = {"scope_user_search_to_org": True}
|
||||
|
||||
async def mock_find_unique_settings(*args, **kwargs):
|
||||
return mock_settings_row
|
||||
|
||||
mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())
|
||||
|
||||
Reference in New Issue
Block a user