diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index daccdd2829d..e33f87494f4 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -140,6 +140,12 @@ from litellm.redis.stores.spend_counters import ( team_member_counter, ) from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.native_auth_repository import ( + EndUserRegistryRow, + ModelAccessGroupRow, + TagRegistryRow, + native_auth_repository, +) from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import DatabaseClient, RowT_co @@ -1494,8 +1500,13 @@ async def get_team_member_default_budget( return cached_budget try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( - where={"budget_id": budget_id} + native: Final = native_auth_repository(prisma_client) + budget_record: Final = ( + await native.one("budget", [budget_id], LiteLLM_BudgetTable) + if native is not None + else await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( + where={"budget_id": budget_id} + ) ) except NativeAuthFailure: raise @@ -1509,7 +1520,9 @@ async def get_team_member_default_budget( verbose_proxy_logger.warning("Team-default member budget not found in database: %s", budget_id) return None - budget: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) + budget: Final = LiteLLM_BudgetTable.model_validate( + budget_record.model_dump() if isinstance(budget_record, BaseModel) else budget_record.dict() + ) await user_api_key_cache.async_set_cache( key=cache_key, value=budget, @@ -1721,9 +1734,16 @@ async def load_end_user_restricted_registry( """The set of end-user ids whose ``LiteLLM_EndUserTable`` row carries a restriction.""" async def fetch_ids() -> tuple[str, ...]: - restricted_rows: Final = await _end_user_table(EndUserRepository(prisma_client)).find_many( - where=_restricted_end_user_where(), - take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1, + native: Final = native_auth_repository(prisma_client) + restricted_rows: Final = ( + await native.many( + "end_user_restricted_registry", [str(END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1)], EndUserRegistryRow + ) + if native is not None + else await _end_user_table(EndUserRepository(prisma_client)).find_many( + where=_restricted_end_user_where(), + take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1, + ) ) return tuple(row.user_id for row in restricted_rows) @@ -1820,16 +1840,23 @@ async def get_end_user_object( # Fetch from database try: - response: Final = await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique( - where={"user_id": end_user_id}, - include={"litellm_budget_table": True, "object_permission": True}, + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("end_user", [end_user_id], LiteLLM_EndUserTable) + if native is not None + else await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique( + where={"user_id": end_user_id}, + include={"litellm_budget_table": True, "object_permission": True}, + ) ) if response is None: raise Exception end_user_row: Final = await _apply_default_budget_to_end_user( - end_user_obj=LiteLLM_EndUserTable.model_validate(response.dict()), + end_user_obj=LiteLLM_EndUserTable.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ), prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, @@ -1911,8 +1938,13 @@ async def _load_tag_registry( """The set of tag names that have a row in ``LiteLLM_TagTable``.""" async def fetch_ids() -> tuple[str, ...]: - registry_rows: Final = await _tag_table(TagRepository(prisma_client)).find_many( - take=TAG_REGISTRY_MAX_SIZE + 1, + native: Final = native_auth_repository(prisma_client) + registry_rows: Final = ( + await native.many("tag_registry", [str(TAG_REGISTRY_MAX_SIZE + 1)], TagRegistryRow) + if native is not None + else await _tag_table(TagRepository(prisma_client)).find_many( + take=TAG_REGISTRY_MAX_SIZE + 1, + ) ) return tuple(row.tag_name for row in registry_rows) @@ -1933,9 +1965,16 @@ async def load_model_access_group_registry( """The set of model access group names that have a row in ``LiteLLM_ModelAccessGroupBudgetTable``.""" async def fetch_ids() -> tuple[str, ...]: - registry_rows: Final = await _model_access_group_budget_table( - ModelAccessGroupBudgetRepository(prisma_client) - ).find_many(take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1) + native: Final = native_auth_repository(prisma_client) + registry_rows: Final = ( + await native.many( + "model_access_group_registry", [str(MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1)], ModelAccessGroupRow + ) + if native is not None + else await _model_access_group_budget_table(ModelAccessGroupBudgetRepository(prisma_client)).find_many( + take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1 + ) + ) return tuple(row.access_group_name for row in registry_rows) return await _load_bounded_registry( @@ -1967,11 +2006,14 @@ async def _fetch_uncached_model_access_group_budgets( return () try: - db_rows: Final = await _model_access_group_budget_table( - ModelAccessGroupBudgetRepository(prisma_client) - ).find_many( - where={"access_group_name": {"in": list(uncached_groups)}}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(prisma_client) + db_rows: Final = ( + await native.many("model_access_group", uncached_groups, ModelAccessGroupRow) + if native is not None + else await _model_access_group_budget_table(ModelAccessGroupBudgetRepository(prisma_client)).find_many( + where={"access_group_name": {"in": list(uncached_groups)}}, + include={"litellm_budget_table": True}, + ) ) fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows) for fetched_name, fetched_obj in fetched: @@ -2053,11 +2095,24 @@ async def _fetch_uncached_tags( return () try: - db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many( - where={"tag_name": {"in": list(tags_to_fetch)}}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(prisma_client) + db_tags: Final = ( + await native.many("tag", tags_to_fetch, LiteLLM_TagTable) + if native is not None + else await _tag_table(TagRepository(prisma_client)).find_many( + where={"tag_name": {"in": list(tags_to_fetch)}}, + include={"litellm_budget_table": True}, + ) + ) + fetched: Final = tuple( + ( + db_tag.tag_name, + LiteLLM_TagTable.model_validate( + db_tag.model_dump() if isinstance(db_tag, BaseModel) else db_tag.dict() + ), + ) + for db_tag in db_tags ) - fetched: Final = tuple((db_tag.tag_name, LiteLLM_TagTable.model_validate(db_tag.dict())) for db_tag in db_tags) for fetched_name, fetched_obj in fetched: await user_api_key_cache.async_set_cache( key=tag_cache_key(fetched_name), @@ -2209,13 +2264,24 @@ async def _fetch_team_membership_from_db( use_writer: bool = False, ) -> LiteLLM_TeamMembership | None: _ = parent_otel_span, proxy_logging_obj - response: Final = await _dictable_table( - TeamMembershipRepository(prisma_client, use_writer=use_writer), "team_membership" - ).find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("member", [user_id, team_id], LiteLLM_TeamMembership, writer=use_writer) + if native is not None + else await _dictable_table( + TeamMembershipRepository(prisma_client, use_writer=use_writer), "team_membership" + ).find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + include={"litellm_budget_table": True}, + ) + ) + membership: Final = ( + None + if response is None + else LiteLLM_TeamMembership.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ) ) - membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) if membership is None: await user_api_key_cache.async_set_cache( @@ -2427,18 +2493,27 @@ async def _get_fuzzy_user_object( """ response = None + native: Final = native_auth_repository(prisma_client) if sso_user_id is not None: - response = await _user_table(UserRepository(prisma_client)).find_unique( - where={"sso_user_id": sso_user_id}, - include={"organization_memberships": True}, + response = ( + await native.one("user_sso", [sso_user_id], LiteLLM_UserTable) + if native is not None + else await _user_table(UserRepository(prisma_client)).find_unique( + where={"sso_user_id": sso_user_id}, + include={"organization_memberships": True}, + ) ) if response is None and user_email is not None: # Use case-insensitive query to handle emails with different casing # This matches the pattern used in _check_duplicate_user_email - response = await _user_table(UserRepository(prisma_client)).find_first( - where={"user_email": {"equals": user_email, "mode": "insensitive"}}, - include={"organization_memberships": True}, + response = ( + await native.one("user_email", [user_email], LiteLLM_UserTable) + if native is not None + else await _user_table(UserRepository(prisma_client)).find_first( + where={"user_email": {"equals": user_email, "mode": "insensitive"}}, + include={"organization_memberships": True}, + ) ) if response is not None and sso_user_id is not None: # update sso_user_id @@ -2528,8 +2603,13 @@ async def get_user_object( ) if should_check_db: - response = await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).find_unique( - where={"user_id": user_id}, include={"organization_memberships": True} + native: Final = native_auth_repository(prisma_client) + response = ( + await native.one("user", [user_id], LiteLLM_UserTable, writer=bool(check_db_only)) + if native is not None + else await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).find_unique( + where={"user_id": user_id}, include={"organization_memberships": True} + ) ) if response is None: @@ -3074,8 +3154,13 @@ class AccessGroupNotFoundError(HTTPException): async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None, *, use_writer: bool = False ) -> "_PrismaTeamRow | None": - response = await _team_table(TeamRepository(prisma_client, use_writer=use_writer)).find_unique( - where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS + native: Final = native_auth_repository(prisma_client) + response = ( + await native.one("team", [team_id], LiteLLM_TeamTable, writer=use_writer) + if native is not None + else await _team_table(TeamRepository(prisma_client, use_writer=use_writer)).find_unique( + where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS + ) ) if response is None and team_id_upsert: @@ -3321,14 +3406,21 @@ async def get_access_object( # Not in cache - fetch from DB try: - response: Final = await _dictable_table( - AccessGroupRepository(prisma_client, use_writer=check_db_only), "access_group" - ).find_unique(where={"access_group_id": access_group_id}) + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("access_group", [access_group_id], LiteLLM_AccessGroupTable, writer=check_db_only) + if native is not None + else await _dictable_table( + AccessGroupRepository(prisma_client, use_writer=check_db_only), "access_group" + ).find_unique(where={"access_group_id": access_group_id}) + ) if response is None: raise AccessGroupNotFoundError(access_group_id=access_group_id) - _response: Final = LiteLLM_AccessGroupTable.model_validate(response.dict()) + _response: Final = LiteLLM_AccessGroupTable.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ) # Save to cache await cache_access_object( @@ -3400,8 +3492,13 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams: Final = await _team_table(TeamRepository(prisma_client)).find_many( - where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS + native: Final = native_auth_repository(prisma_client) + teams: Final = ( + await native.many("team_alias", [team_alias], LiteLLM_TeamTable) + if native is not None + else await _team_table(TeamRepository(prisma_client)).find_many( + where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS + ) ) if not teams: @@ -3512,8 +3609,13 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many( - where={"organization_alias": org_alias} + native: Final = native_auth_repository(prisma_client) + orgs = ( + await native.many("org_alias", [org_alias], LiteLLM_OrganizationTable) + if native is not None + else await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many( + where={"organization_alias": org_alias} + ) ) if not orgs: @@ -3735,6 +3837,9 @@ async def _fetch_key_object_from_db_unbounded( *, check_db_only: bool = False, ) -> BaseModel | None: + native: Final = native_auth_repository(prisma_client) + if native is not None: + return await native.key(hashed_token, writer=check_db_only) fetch: Final = partial(prisma_client.get_data, use_writer=True) if check_db_only else prisma_client.get_data async with db_lookup_gate.current(): try: @@ -3978,16 +4083,25 @@ async def get_object_permission( # else, check db try: - response: Final = await _dictable_table( - ObjectPermissionRepository(prisma_client, use_writer=check_db_only), "object_permission" - ).find_unique(where={"object_permission_id": object_permission_id}) + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one( + "object_permission", [object_permission_id], LiteLLM_ObjectPermissionTable, writer=check_db_only + ) + if native is not None + else await _dictable_table( + ObjectPermissionRepository(prisma_client, use_writer=check_db_only), "object_permission" + ).find_unique(where={"object_permission_id": object_permission_id}) + ) if response is None: if check_db_only: raise HTTPException(status_code=403, detail="Referenced object permission does not exist") return None - _perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate(response.dict()) + _perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate( + response.model_dump() if isinstance(response, BaseModel) else response.dict() + ) await user_api_key_cache.async_set_cache( key=key, value=_perm_obj, @@ -4119,8 +4233,11 @@ async def get_org_object( if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique( - **query_kwargs + native: Final = native_auth_repository(prisma_client) + response: Final = ( + await native.one("org", [org_id], LiteLLM_OrganizationTable) + if native is not None + else await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique(**query_kwargs) ) if response is None: @@ -5456,9 +5573,14 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row: Final = await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique( - where={"project_id": project_id}, - include={"litellm_budget_table": True}, + native: Final = native_auth_repository(prisma_client) + project_row: Final = ( + await native.one("project", [project_id], LiteLLM_ProjectTableCachedObj) + if native is not None + else await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique( + where={"project_id": project_id}, + include={"litellm_budget_table": True}, + ) ) if project_row is None: return None diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index 6257d461cf4..c1463dd6909 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -4,6 +4,7 @@ the readers and the fallback, so enforcement never depends on this running.""" from __future__ import annotations +import json import time from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass @@ -12,6 +13,7 @@ from typing import Final, Literal, TypeAlias from pydantic import BaseModel, TypeAdapter, ValidationError +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._internal_context import service_target from litellm._logging import verbose_proxy_logger from litellm.caching.redis_batch import Batch @@ -32,6 +34,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.db.db_span import db_span from litellm.proxy.utils import PrismaClient from litellm.redis.stores.auth_cache import AuthCacheStore +from litellm.repositories.native_auth_repository import native_auth_repository _RowKind: TypeAlias = Literal["user_row", "team_row", "membership_row", "organization_row", "project_row"] @@ -235,7 +238,7 @@ async def _fill_from_redis( def _validate_row( - row_value: object, model_type: type[BaseModel], row: _RowKind, refreshed_at: float + row_value: object, model_type: type[BaseModel], row: _RowKind, refreshed_at: float, *, native: bool = False ) -> BaseModel | None: if row_value is None: return None @@ -246,6 +249,8 @@ def _validate_row( return model_type.model_validate(stamped) return model_type.model_validate(columns) except ValidationError as e: + if native: + raise NativeAuthRuntimeError(f"Invalid native auth {row} projection") from e verbose_proxy_logger.warning("auth prefetch: %s did not validate as %s: %s", row, model_type.__name__, e) return None @@ -253,6 +258,20 @@ def _validate_row( async def _fetch_rows( refs: AuthObjectRefs, kinds: frozenset[_RowKind], prisma_client: PrismaClient ) -> Mapping[str, object]: + native: Final = native_auth_repository(prisma_client) + if native is not None: + refs_json: Final = json.dumps( + [ + refs.user_id if "user_row" in kinds else None, + refs.team_id if kinds & _TEAM_BOUND_ROWS else None, + refs.membership_user_id if "membership_row" in kinds else None, + refs.organization_id if "organization_row" in kinds else None, + refs.project_id if "project_row" in kinds else None, + ] + ) + async with db_span("prefetch_auth_objects", AUTH_OBJECTS_TARGET): + rows: Final = await native.read("prefetch", [refs_json]) + return rows[0] if rows else _NO_ROWS async with db_span("prefetch_auth_objects", AUTH_OBJECTS_TARGET): row: Final[object] = await prisma_client.db.query_first( # pyright: ignore[reportAny] # prisma types query_first as Any _SQL, @@ -289,8 +308,12 @@ async def _fill_from_db( ) rows: Final = await _fetch_rows(refs, frozenset(model_for), prisma_client) refreshed_at: Final = time.time() + native: Final = native_auth_repository(prisma_client) is not None objects: Final[Mapping[_RowKind, BaseModel | None]] = MappingProxyType( - {row: _validate_row(rows.get(row), model_type, row, refreshed_at) for row, model_type in model_for.items()} + { + row: _validate_row(rows.get(row), model_type, row, refreshed_at, native=native) + for row, model_type in model_for.items() + } ) writes: Final = tuple((entry, value) for entry in entries if (value := objects[entry.row]) is not None) if writes: @@ -313,6 +336,8 @@ async def prefetch_auth_objects( if prisma_client is None: return await _fill_from_db(refs, _missing_in_memory(missing, user_api_key_cache), user_api_key_cache, prisma_client) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # warm-up only; the getters enforce and fail closed on their own verbose_proxy_logger.warning("auth prefetch skipped, falling back to per-object lookups: %s", e) diff --git a/litellm/proxy/auth/end_user_context_io.py b/litellm/proxy/auth/end_user_context_io.py index 83cf9784886..e39ad797bea 100644 --- a/litellm/proxy/auth/end_user_context_io.py +++ b/litellm/proxy/auth/end_user_context_io.py @@ -9,6 +9,7 @@ from litellm.proxy._types import LiteLLM_BudgetTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.native_auth_repository import native_auth_repository @dataclass(frozen=True, slots=True) @@ -29,14 +30,22 @@ async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTa if cached is not None: return cached try: - table: Final = auth_checks._dictable_table(BudgetRepository(io.prisma), "budget") # pyright: ignore[reportPrivateUsage] # shared raw repository projection - record: Final = await table.find_unique(where={"budget_id": budget_id}) + native: Final = native_auth_repository(io.prisma) + record: Final = ( + await native.one("budget", [budget_id], LiteLLM_BudgetTable) + if native is not None + else await auth_checks._dictable_table(BudgetRepository(io.prisma), "budget").find_unique( + where={"budget_id": budget_id} + ) # pyright: ignore[reportPrivateUsage] # shared raw repository projection + ) if record is None: verbose_proxy_logger.warning( "Default end user budget not found in database: %s", budget_id.replace("\r", "").replace("\n", "") ) return None - row: Final = LiteLLM_BudgetTable.model_validate(record.dict()) + row: Final = LiteLLM_BudgetTable.model_validate( + record.model_dump() if isinstance(record, LiteLLM_BudgetTable) else record.dict() + ) await io.cache.async_set_cache( key=key, value=row, model_type=LiteLLM_BudgetTable, ttl=get_management_object_ttl(io.cache) ) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 7f2aa33246d..2646b28e8a1 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -14,14 +14,16 @@ from types import MappingProxyType from typing import Final, TypeVar import litellm +from litellm._auth.failure import NativeAuthFailure from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping -from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialModelNames, UserAPIKeyAuth from litellm.proxy.common_utils.model_listing_utils import CallerAliases, alias_map, dispatched_model from litellm.proxy.common_utils.token_count_settings import cost_map_entry +from litellm.repositories.native_auth_repository import native_auth_repository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.responses.main import strip_responses_routing_prefix from litellm.router import Router @@ -132,12 +134,21 @@ async def get_mcp_server_ids( # Make a direct SQL query to get just the mcp_servers try: - result: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( - where={"object_permission_id": user_api_key_dict.object_permission_id}, + native: Final = native_auth_repository(prisma_client) + result: Final = ( + await native.one( + "object_permission", [user_api_key_dict.object_permission_id], LiteLLM_ObjectPermissionTable + ) + if native is not None + else await ObjectPermissionRepository(prisma_client).table.find_unique( + where={"object_permission_id": user_api_key_dict.object_permission_id}, + ) ) if result and result.mcp_servers: return result.mcp_servers return [] + except NativeAuthFailure: + raise except Exception: return [] diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4ed3ac30be7..a944bd6b513 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -30,6 +30,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import ( record_spend_counter_value, ) from litellm.redis.stores.spend_counters import END_USER_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendCounterStore +from litellm.repositories.native_auth_repository import BudgetWindowRow, SpendRow, native_auth_repository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( @@ -134,6 +135,20 @@ class SpendCounterReseed: # or tag names that merely contain ":window:". if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + prefix, separator, suffix = counter_key.partition("spend:") + kind, delimiter, entity_id = suffix.partition(":") + if ( + prefix + or not separator + or not delimiter + or kind not in {"key", "team", "team_member", "user", "org", "project"} + ): + return None + if kind == "team_member" and ":" not in entity_id: + return None + return await native.spend("member" if kind == "team_member" else kind, entity_id) try: row: Final = await bounded_db_lookup( SpendCounterReseed._counter_row(prisma_client, counter_key), name="spend_counter" @@ -181,6 +196,9 @@ class SpendCounterReseed: async def end_user_from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: if prisma_client is None or not counter_key.startswith(END_USER_COUNTER_PREFIX): return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + return await native.spend("end_user", counter_key[len(END_USER_COUNTER_PREFIX) :]) where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]} try: row: Final = await bounded_db_lookup( @@ -296,6 +314,15 @@ class SpendCounterReseed: if row_entity_type is None: return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + window: Final = await native.one( + "budget_window", [entity_type, entity_id, window_duration], BudgetWindowRow + ) + if window is None or _as_utc(window.window_start) < _as_utc(expected_window_start): + return None + return float(window.spend or 0.0) + try: row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique( where={ @@ -367,6 +394,12 @@ class SpendCounterReseed: group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) if group_field is None: return None + native: Final = native_auth_repository(prisma_client) + if native is not None: + aggregate: Final = await native.one( + "window_spend", [entity_type, entity_id, window_start.isoformat()], SpendRow + ) + return float(aggregate.spend or 0.0) if aggregate is not None else 0.0 where: Final = { group_field: entity_id, "startTime": {"gte": window_start}, diff --git a/litellm/repositories/native_auth_repository.py b/litellm/repositories/native_auth_repository.py new file mode 100644 index 00000000000..74f3e09ece2 --- /dev/null +++ b/litellm/repositories/native_auth_repository.py @@ -0,0 +1,143 @@ +import time +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime +from inspect import getattr_static +from typing import Final, Protocol, TypeVar + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm.litellm_core_utils.core_helpers import require_attributes +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.team import Member +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded, bounded_db_lookup, db_lookup_gate +from litellm.types.proxy.auth.user_api_key_auth import LiteLLM_VerificationTokenView + +_Model = TypeVar("_Model", bound=BaseModel) +_ROWS: Final = TypeAdapter(list[dict[str, object]]) +_SPEND: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) +_MEMBERS: Final = TypeAdapter(list[dict[str, object]]) + + +class ModelAccessGroupRow(BaseModel): + model_config = ConfigDict(frozen=True) + access_group_name: str + spend: float | None = 0.0 + litellm_budget_table: LiteLLM_BudgetTable | None = None + + +class EndUserRegistryRow(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str + + +class TagRegistryRow(BaseModel): + model_config = ConfigDict(frozen=True) + tag_name: str + + +class SpendRow(BaseModel): + model_config = ConfigDict(frozen=True) + spend: float | None = None + + +class BudgetWindowRow(SpendRow): + window_start: datetime + + +class NativeAuthDeadlineExceeded(NativeAuthFailure, DBLookupDeadlineExceeded): + pass + + +class AuthReadHandle(Protocol): + async def auth_read(self, kind: str, ids: list[str], writer: bool = False) -> object: ... + + async def auth_spend(self, kind: str, id: str) -> object: ... + + +class _NativeAuthClient(Protocol): + @property + def native_auth_database(self) -> AuthReadHandle | None: ... + + +@dataclass(frozen=True, slots=True) +class NativeAuthRepository: + database: AuthReadHandle + + async def key(self, token: str, *, writer: bool = False) -> LiteLLM_VerificationTokenView | None: + rows: Final = await self.read("key", [token], writer=writer) + if not rows: + return None + try: + row: Final = rows[0] + members: Final = _MEMBERS.validate_python(row.get("team_members_with_roles") or []) + matches: Final = tuple( + Member.model_validate(member) + for member in members + if member.get("user_id") is not None and member.get("user_id") == row.get("user_id") + ) + projection: Final = LiteLLM_VerificationTokenView.model_validate( + { + **row, + "team_models": row.get("team_models") if row.get("team_models") is not None else [], + "team_blocked": row.get("team_blocked") if row.get("team_blocked") is not None else False, + "team_member": matches[-1] if matches else None, + "last_refreshed_at": time.time(), + } + ) + return ( + projection.model_copy(update={"expires": projection.expires.isoformat()}) + if isinstance(projection.expires, datetime) + else projection + ) + except Exception as error: # noqa: BLE001 # preserve terminal native projection errors + raise NativeAuthRuntimeError("Invalid native auth key projection") from error + + async def read(self, kind: str, ids: Sequence[str], writer: bool = False) -> list[dict[str, object]]: + try: + value: Final = await bounded_db_lookup(self._read(kind, ids, writer), name=kind) + return _ROWS.validate_python(value) + except NativeAuthFailure: + raise + except DBLookupDeadlineExceeded as error: + raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error + except Exception as error: # noqa: BLE001 # native errors must not enter legacy fallback paths + raise NativeAuthRuntimeError(f"Native auth {kind} read failed") from error + + async def _read(self, kind: str, ids: Sequence[str], writer: bool) -> object: + async with db_lookup_gate.current(): + return await self.database.auth_read(kind, list(ids), writer=writer) + + async def one(self, kind: str, ids: Sequence[str], model: type[_Model], *, writer: bool = False) -> _Model | None: + rows: Final = await self.many(kind, ids, model, writer=writer) + return rows[0] if rows else None + + async def many(self, kind: str, ids: Sequence[str], model: type[_Model], *, writer: bool = False) -> list[_Model]: + rows: Final = await self.read(kind, ids, writer=writer) + try: + return [model.model_validate(row) for row in rows] + except Exception as error: # noqa: BLE001 # invalid native projections are terminal + raise NativeAuthRuntimeError(f"Invalid native auth {kind} projection") from error + + async def spend(self, kind: str, entity_id: str) -> float | None: + try: + value: Final = await bounded_db_lookup(self._spend(kind, entity_id), name="spend_counter") + return _SPEND.validate_python(value) + except NativeAuthFailure: + raise + except DBLookupDeadlineExceeded as error: + raise NativeAuthDeadlineExceeded(error.lookup, error.deadline_seconds) from error + except Exception as error: # noqa: BLE001 # a failed native read is not an authoritative miss + raise NativeAuthRuntimeError(f"Native auth {kind} spend read failed") from error + + async def _spend(self, kind: str, entity_id: str) -> object: + async with db_lookup_gate.current(): + return await self.database.auth_spend(kind, entity_id) + + +def native_auth_repository(client: object) -> NativeAuthRepository | None: + if getattr_static(client, "native_auth_database", None) is None: + return None + database: Final = require_attributes(client, _NativeAuthClient, "native_auth_database").native_auth_database + return NativeAuthRepository(database) if database is not None else None diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 3038e62d219..8f4404566fd 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -632,6 +632,37 @@ async def test_get_key_object_should_reconnect_once_on_db_connection_error(): ) +@pytest.mark.asyncio +async def test_native_key_lookup_populates_cache_without_prisma_io(): + native = SimpleNamespace(auth_read=AsyncMock(return_value=[{"token": "native-hash", "team_models": None, "team_blocked": None}])) + prisma = MagicMock(native_auth_database=native) + prisma.get_data = AsyncMock(side_effect=AssertionError("legacy read")) + prisma.attempt_db_reconnect = AsyncMock(side_effect=AssertionError("legacy reconnect")) + cache = UserApiKeyCache() + first = await get_key_object("native-hash", prisma, cache) + second = await get_key_object("native-hash", prisma, cache) + assert first.token == second.token == "native-hash" + assert first.team_models == second.team_models == [] + assert native.auth_read.await_count == 1 + prisma.get_data.assert_not_awaited() + prisma.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_native_key_transport_error_never_enters_prisma_reconnect(): + from litellm._auth.failure import NativeAuthFailure + + native = SimpleNamespace(auth_read=AsyncMock(side_effect=httpx.ConnectError("db offline"))) + prisma = MagicMock(native_auth_database=native) + prisma.get_data = AsyncMock() + prisma.attempt_db_reconnect = AsyncMock() + with pytest.raises(NativeAuthFailure): + await fetch_key_object_from_db_with_reconnect("hash", prisma, None, None, check_db_only=True) + native.auth_read.assert_awaited_once_with("key", ["hash"], writer=True) + prisma.get_data.assert_not_awaited() + prisma.attempt_db_reconnect.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_error(): mock_prisma_client = MagicMock() diff --git a/tests/unit/proxy/auth/test_auth_object_prefetch.py b/tests/unit/proxy/auth/test_auth_object_prefetch.py index a80c637d3df..d6c5d4559d4 100644 --- a/tests/unit/proxy/auth/test_auth_object_prefetch.py +++ b/tests/unit/proxy/auth/test_auth_object_prefetch.py @@ -1,11 +1,14 @@ """Counts the Redis round trips and DB queries auth object reads cost per cache regime, and checks that the per-object getters still enforce on their own when the prefetch cannot help.""" +import json +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from litellm._auth.failure import NativeAuthFailure from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy._types import ( LiteLLM_OrganizationTable, @@ -75,6 +78,51 @@ ALL_ROWS = { } +@pytest.mark.asyncio +async def test_native_prefetch_fills_all_getters_with_one_read_and_preserves_optional_ids(): + native = SimpleNamespace(auth_read=AsyncMock(return_value=[ALL_ROWS])) + prisma = _prisma() + prisma.native_auth_database = native + cache = _cache(None) + + await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=cache) + user, team, member, org = await _read_all_through_getters(cache, prisma) + + assert user is not None and user.user_id == USER_ID + assert team.team_id == TEAM_ID + assert member is not None and member.spend == 3.0 + assert org is not None and org.litellm_budget_table.max_budget == 1000.0 + assert native.auth_read.await_count == 1 + kind, ids = native.auth_read.call_args.args + assert kind == "prefetch" + assert len(ids) == 1 and json.loads(ids[0]) == [USER_ID, TEAM_ID, USER_ID, ORG_ID, None] + assert prisma.db.mock_calls == [] + + await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=cache) + assert native.auth_read.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rows", [[{"user_row": {"spend": "bad"}}], {"user_row": USER_ROW}]) +async def test_native_prefetch_invalid_projection_is_terminal_without_legacy_query(rows): + prisma = _prisma() + prisma.native_auth_database = SimpleNamespace(auth_read=AsyncMock(return_value=rows)) + with pytest.raises(NativeAuthFailure): + await prefetch_auth_objects( + refs=AuthObjectRefs(user_id=USER_ID), prisma_client=prisma, user_api_key_cache=_cache(None) + ) + assert prisma.db.mock_calls == [] + + +@pytest.mark.asyncio +async def test_native_prefetch_transport_failure_is_not_best_effort(): + prisma = _prisma() + prisma.native_auth_database = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("offline"))) + with pytest.raises(NativeAuthFailure): + await prefetch_auth_objects(refs=_refs(), prisma_client=prisma, user_api_key_cache=_cache(None)) + assert prisma.db.mock_calls == [] + + def _prisma(rows: dict[str, object] | None = ALL_ROWS) -> MagicMock: prisma = MagicMock(name="prisma_client") prisma.db.query_first = AsyncMock(return_value=rows) diff --git a/tests/unit/proxy/db/test_spend_counter_reseed.py b/tests/unit/proxy/db/test_spend_counter_reseed.py index 7314bfbb7c5..62ceff452e7 100644 --- a/tests/unit/proxy/db/test_spend_counter_reseed.py +++ b/tests/unit/proxy/db/test_spend_counter_reseed.py @@ -16,6 +16,7 @@ from unittest.mock import AsyncMock import pytest +from litellm._auth.failure import NativeAuthFailure from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY @@ -27,6 +28,69 @@ from litellm.redis.stores.spend_counters import SpendCounterStore WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("counter", "kind", "entity_id"), + [ + ("spend:key:k", "key", "k"), + ("spend:user:u", "user", "u"), + ("spend:team:t", "team", "t"), + ("spend:org:o", "org", "o"), + ("spend:project:p", "project", "p"), + ("spend:team_member:u:with:colon:t", "member", "u:with:colon:t"), + ], +) +async def test_native_reseed_reads_authoritative_zero_without_a_prisma_handle(counter, kind, entity_id): + native = SimpleNamespace(auth_spend=AsyncMock(return_value=0.0)) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.from_db(prisma, counter) == 0.0 + native.auth_spend.assert_awaited_once_with(kind, entity_id) + + +@pytest.mark.asyncio +async def test_native_reseed_distinguishes_missing_from_failure(): + native = SimpleNamespace(auth_spend=AsyncMock(side_effect=[None, RuntimeError("offline")])) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.from_db(prisma, "spend:user:u") is None + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.from_db(prisma, "spend:user:u") + assert native.auth_spend.await_count == 2 + + +@pytest.mark.asyncio +async def test_native_end_user_floor_preserves_zero(): + native = SimpleNamespace(auth_spend=AsyncMock(return_value=0.0)) + assert ( + await SpendCounterReseed.end_user_from_db(SimpleNamespace(native_auth_database=native), "spend:end_user:u") + == 0.0 + ) + native.auth_spend.assert_awaited_once_with("end_user", "u") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("row", [None, {"spend": 99.0, "window_start": WINDOW_START - timedelta(days=1)}]) +async def test_native_stale_or_missing_window_uses_native_spend_aggregate(row): + native = SimpleNamespace(auth_read=AsyncMock(side_effect=[[row] if row else [], [{"spend": 4.5}]])) + value = await SpendCounterReseed.window_from_db( + SimpleNamespace(native_auth_database=native), "Key", "hash", "1d", WINDOW_START + ) + assert value == 4.5 + assert native.auth_read.call_args_list[0].args == ("budget_window", ["Key", "hash", "1d"]) + assert native.auth_read.call_args_list[1].args == ("window_spend", ["Key", "hash", WINDOW_START.isoformat()]) + + +@pytest.mark.asyncio +async def test_native_current_window_does_not_range_scan_and_failure_does_not_fallback(): + native = SimpleNamespace( + auth_read=AsyncMock(side_effect=[[{"spend": 0.0, "window_start": WINDOW_START}], RuntimeError("offline")]) + ) + prisma = SimpleNamespace(native_auth_database=native) + assert await SpendCounterReseed.window_from_db(prisma, "Key", "hash", "1d", WINDOW_START) == 0.0 + with pytest.raises(NativeAuthFailure): + await SpendCounterReseed.window_from_db(prisma, "Key", "hash", "1d", WINDOW_START) + assert native.auth_read.await_count == 2 + + class _FakeFindUniqueTable: def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: self._row = row diff --git a/tests/unit/repositories/test_native_auth_repository.py b/tests/unit/repositories/test_native_auth_repository.py new file mode 100644 index 00000000000..686135e123e --- /dev/null +++ b/tests/unit/repositories/test_native_auth_repository.py @@ -0,0 +1,104 @@ +import asyncio +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.repositories.native_auth_repository import NativeAuthRepository, native_auth_repository + + +@pytest.mark.asyncio +async def test_native_key_preserves_joined_budget_and_membership_projection(): + raw = { + "token": "token-hash", + "user_id": "user", + "organization_id": "key-org", + "org_id": "team-org", + "team_models": None, + "team_blocked": None, + "max_budget": None, + "litellm_budget_table_max_budget": 25.0, + "model_max_budget": {}, + "litellm_budget_table_model_max_budget": {"model": {"max_budget": 1.0}}, + "team_model_aliases": {"alias": "model"}, + "team_members_with_roles": [{"user_id": "user", "role": "user"}, {"user_id": "user", "role": "admin"}], + "expires": datetime(2027, 1, 1, tzinfo=timezone.utc), + } + handle = SimpleNamespace(auth_read=AsyncMock(return_value=[raw])) + row = await NativeAuthRepository(handle).key("token-hash", writer=True) + + assert row is not None + assert (row.org_id, row.team_models, row.team_blocked, row.max_budget) == ("key-org", [], False, 25.0) + assert row.model_max_budget == {"model": {"max_budget": 1.0}} + assert row.team_model_aliases == {"alias": "model"} + assert row.team_member is not None and row.team_member.role == "admin" + assert row.expires == "2027-01-01T00:00:00+00:00" + assert row.last_refreshed_at is not None + assert raw["max_budget"] is None and raw["team_models"] is None + handle.auth_read.assert_awaited_once_with("key", ["token-hash"], writer=True) + + +@pytest.mark.asyncio +async def test_native_rows_keep_null_limits_and_make_request_local_models(): + raw = {"user_id": "u", "team_id": "t", "spend": None, "litellm_budget_table": {"max_budget": None}} + handle = SimpleNamespace(auth_read=AsyncMock(return_value=[raw])) + repo = NativeAuthRepository(handle) + first = await repo.one("member", ["u", "t"], LiteLLM_TeamMembership) + second = await repo.one("member", ["u", "t"], LiteLLM_TeamMembership) + + assert first is not None and second is not None + assert first.spend is None and second.spend is None + assert first.litellm_budget_table is not None and second.litellm_budget_table is not None + first.litellm_budget_table.max_budget = 3.0 + assert second.litellm_budget_table.max_budget is None + assert raw["litellm_budget_table"] == {"max_budget": None} + + +@pytest.mark.asyncio +async def test_native_missing_row_is_none_and_zero_budget_is_not_missing(): + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=[[], [{"budget_id": "b", "max_budget": 0.0}]])) + repo = NativeAuthRepository(handle) + assert await repo.one("budget", ["missing"], LiteLLM_BudgetTable) is None + found = await repo.one("budget", ["b"], LiteLLM_BudgetTable) + assert found is not None and found.max_budget == 0.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("answer", [None, {}, ["bad"], [{"max_budget": "bad"}]]) +async def test_native_invalid_rows_are_terminal(answer): + handle = SimpleNamespace(auth_read=AsyncMock(return_value=answer)) + with pytest.raises(NativeAuthFailure): + await NativeAuthRepository(handle).one("budget", ["b"], LiteLLM_BudgetTable) + assert handle.auth_read.await_count == 1 + + +@pytest.mark.asyncio +async def test_native_failure_and_cancellation_are_not_clean_misses(): + handle = SimpleNamespace(auth_read=AsyncMock(side_effect=RuntimeError("connection closed"))) + with pytest.raises(NativeAuthFailure, match="read failed"): + await NativeAuthRepository(handle).key("hash") + cancelled = SimpleNamespace(auth_read=AsyncMock(side_effect=asyncio.CancelledError)) + with pytest.raises(asyncio.CancelledError): + await NativeAuthRepository(cancelled).key("hash") + + +@pytest.mark.asyncio +async def test_native_spend_preserves_missing_and_authoritative_zero(): + handle = SimpleNamespace(auth_spend=AsyncMock(side_effect=[None, 0.0, RuntimeError("offline")])) + repo = NativeAuthRepository(handle) + assert await repo.spend("user", "u") is None + assert await repo.spend("user", "u") == 0.0 + with pytest.raises(NativeAuthFailure): + await repo.spend("user", "u") + + +def test_repository_selection_requires_an_explicit_handle(): + assert native_auth_repository(SimpleNamespace(native_auth_database=None)) is None + assert native_auth_repository(MagicMock()) is None + handle = SimpleNamespace(auth_read=AsyncMock()) + repo = native_auth_repository(SimpleNamespace(native_auth_database=handle)) + assert repo is not None and repo.database is handle