scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

refactor(auth): share key cache miss loading

Browse source
This commit is contained in:
Pin Zhu 2026-09-28 19:14:53 +08:00 • committed by GitHub
parent 9a38646f35
commit d5b2c586d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 6 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
litellm/proxy/auth/resolvers/store.py
View file

@ -3,15 +3,10 @@ from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
_cache_key_object,
_copy_user_api_key_auth_for_cache,
_fetch_key_object_from_db_with_reconnect,
get_object_permission,
_load_key_object_on_cache_miss,
)
from litellm.proxy.auth.auth_method import AuthMethod
from litellm.proxy.auth.network import NetworkContext
@ -34,8 +29,8 @@ from litellm.proxy.auth.resolvers.models import (
from litellm.proxy.auth.roles import TeamRole, map_role, team_role
if TYPE_CHECKING:
from litellm.caching.caching import DualCache
from litellm.integrations.opentelemetry import Span
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -59,7 +54,7 @@ class IdentityStore:
def __init__(
self,
prisma_client: PrismaClient | None,
cache: DualCache,
cache: UserApiKeyCache,
*,
parent_otel_span: Span | None = None,
proxy_logging_obj: ProxyLogging | None = None,
@ -110,39 +105,15 @@ class IdentityStore:
if self._check_cache_only:
raise KeyNotInCacheError(hashed_token)
from_db: Final[BaseModel | None] = await _fetch_key_object_from_db_with_reconnect(
key: Final = await _load_key_object_on_cache_miss(
hashed_token=hashed_token,
prisma_client=self._prisma,
user_api_key_cache=self._cache,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
)
if from_db is None:
if key is None:
raise KeyNotFoundError(hashed_token)
key: Final = UserAPIKeyAuth.model_validate(from_db.model_dump(exclude_none=True))
if key.object_permission_id and not key.object_permission:
try:
key.object_permission = await get_object_permission(
object_permission_id=key.object_permission_id,
prisma_client=self._prisma,
user_api_key_cache=self._cache,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
)
except Exception as e:
verbose_proxy_logger.debug(
"Failed to load object_permission for key with object_permission_id=%s: %s",
key.object_permission_id,
e,
)
await _cache_key_object(
hashed_token=hashed_token,
user_api_key_obj=key,
user_api_key_cache=self._cache,
proxy_logging_obj=self._proxy_logging_obj,
)
return key
@staticmethod