Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/auth/resolvers/store.py: 81%
62 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1from __future__ import annotations
3from collections.abc import Sequence
4from typing import TYPE_CHECKING, Final
6from pydantic import BaseModel
8from litellm._logging import verbose_proxy_logger
9from litellm.proxy._types import UserAPIKeyAuth
10from litellm.proxy.auth.auth_checks import (
11 _cache_key_object,
12 _copy_user_api_key_auth_for_cache,
13 _fetch_key_object_from_db_with_reconnect,
14 get_object_permission,
15)
16from litellm.proxy.auth.auth_method import AuthMethod
17from litellm.proxy.auth.network import NetworkContext
18from litellm.proxy.auth.resolvers.exceptions import (
19 KeyNotFoundError,
20 KeyNotInCacheError,
21 NoDatabaseConnectionError,
22 PrincipalMissingSourceKeyError,
23)
24from litellm.proxy.auth.resolvers.models import (
25 CredentialRef,
26 EndUserIdentity,
27 OrganizationIdentity,
28 Principal,
29 PrincipalType,
30 ProjectIdentity,
31 TeamIdentity,
32 UserIdentity,
33)
34from litellm.proxy.auth.roles import TeamRole, map_role, team_role
36if TYPE_CHECKING: 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true
37 from litellm.caching.caching import DualCache
38 from litellm.integrations.opentelemetry import Span
39 from litellm.proxy.utils import PrismaClient, ProxyLogging
42class IdentityStore:
43 """The auth flow's resolver: one combined_view lookup, projected into a Principal.
45 ``resolve`` does the lookup (cache, then DB via the shared lower-level helpers,
46 then write-back) and returns the per-caller Principal. The Principal carries the
47 source key object so ``key_from_principal`` can hand it back to the parts of the
48 request flow that still consume ``UserAPIKeyAuth`` (budget, rate limits, policy);
49 that carrier is a stopgap until those consumers read identity off the Principal.
50 The Prisma client, key cache, the request's tracing span / logging sink, and
51 whether this store may only read the cache are injected so the composition root
52 can build the store once the proxy DB is connected; the span and logging sink
53 are infra the DB call is instrumented with and ``check_cache_only`` is a store
54 mode, none of them inputs to resolving identity. ``auth_checks.get_key_object``
55 stays as the legacy entrypoint for its other callers until they migrate onto
56 this store.
57 """
59 def __init__(
60 self,
61 prisma_client: PrismaClient | None,
62 cache: DualCache,
63 *,
64 parent_otel_span: Span | None = None,
65 proxy_logging_obj: ProxyLogging | None = None,
66 check_cache_only: bool = False,
67 ) -> None:
68 self._prisma = prisma_client
69 self._cache = cache
70 self._parent_otel_span = parent_otel_span
71 self._proxy_logging_obj = proxy_logging_obj
72 self._check_cache_only = check_cache_only
74 async def resolve(
75 self,
76 hashed_token: str,
77 *,
78 auth_method: AuthMethod = AuthMethod.API_KEY,
79 network: NetworkContext | None = None,
80 ) -> Principal:
81 key: Final = await self._resolve_key(hashed_token)
82 return self._principal_from_key(
83 key,
84 auth_method=auth_method,
85 network=network,
86 subject_fallback=key.token,
87 credential_ref=CredentialRef(token_id=key.token),
88 )
90 @staticmethod
91 def key_from_principal(principal: Principal) -> UserAPIKeyAuth:
92 """Hand back the resolved key object carried on the Principal.
94 Stopgap for the request flow that still consumes ``UserAPIKeyAuth`` for
95 budget, rate-limit, and policy state. Only Principals produced by
96 ``resolve`` carry a source key.
97 """
98 if principal.source_key is None: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true
99 raise PrincipalMissingSourceKeyError()
100 return principal.source_key
102 async def _resolve_key(self, hashed_token: str) -> UserAPIKeyAuth:
103 if self._prisma is None: 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true
104 raise NoDatabaseConnectionError()
106 cached: Final = await self._cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth)
107 if cached is not None:
108 return _copy_user_api_key_auth_for_cache(user_api_key_obj=cached)
110 if self._check_cache_only:
111 raise KeyNotInCacheError(hashed_token)
113 from_db: Final[BaseModel | None] = await _fetch_key_object_from_db_with_reconnect(
114 hashed_token=hashed_token,
115 prisma_client=self._prisma,
116 parent_otel_span=self._parent_otel_span,
117 proxy_logging_obj=self._proxy_logging_obj,
118 )
119 if from_db is None: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true
120 raise KeyNotFoundError(hashed_token)
122 key: Final = UserAPIKeyAuth.model_validate(from_db.model_dump(exclude_none=True))
124 if key.object_permission_id and not key.object_permission:
125 try:
126 key.object_permission = await get_object_permission(
127 object_permission_id=key.object_permission_id,
128 prisma_client=self._prisma,
129 user_api_key_cache=self._cache,
130 parent_otel_span=self._parent_otel_span,
131 proxy_logging_obj=self._proxy_logging_obj,
132 )
133 except Exception as e:
134 verbose_proxy_logger.debug(
135 "Failed to load object_permission for key with object_permission_id=%s: %s",
136 key.object_permission_id,
137 e,
138 )
140 await _cache_key_object(
141 hashed_token=hashed_token,
142 user_api_key_obj=key,
143 user_api_key_cache=self._cache,
144 proxy_logging_obj=self._proxy_logging_obj,
145 )
146 return key
148 @staticmethod
149 def _principal_from_key(
150 key: UserAPIKeyAuth,
151 *,
152 auth_method: AuthMethod,
153 issuer: str | None = None,
154 subject_fallback: str | None = None,
155 scopes: Sequence[str] = (),
156 credential_ref: CredentialRef | None = None,
157 network: NetworkContext | None = None,
158 ) -> Principal:
159 """Project the identity slice off an already-resolved key object and carry
160 the key on the Principal so ``key_from_principal`` can recover it.
162 Pure: issues no lookup. Both ``resolve`` and the auth seam call this so
163 identity is projected once off whichever key object they already hold.
164 """
165 teams: Final[list[TeamIdentity]] = []
166 if key.team_id is not None: 166 ↛ 167line 166 didn't jump to line 167 because the condition on line 166 was never true
167 role: Final = team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER
168 teams.append(TeamIdentity(id=key.team_id, name=key.team_alias, role=role))
169 organization: Final = (
170 OrganizationIdentity(id=key.org_id, name=key.organization_alias) if key.org_id is not None else None
171 )
172 user: Final = UserIdentity(id=key.user_id, email=key.user_email) if key.user_id is not None else None
173 project = ProjectIdentity(id=key.project_id, name=key.project_alias) if key.project_id is not None else None
174 end_user: Final = EndUserIdentity(id=key.end_user_id) if key.end_user_id is not None else None
175 mapped: Final = map_role(key.user_role)
176 return Principal(
177 principal_type=(PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT),
178 subject=key.user_id or key.key_alias or subject_fallback or "",
179 issuer=issuer,
180 user=user,
181 organization=organization,
182 teams=teams,
183 project=project,
184 end_user=end_user,
185 roles=[mapped] if mapped else [],
186 scopes=list(scopes),
187 auth_method=auth_method,
188 credential_ref=credential_ref or CredentialRef(),
189 network=network or NetworkContext(),
190 source_key=key,
191 )