From 65401138b5856c095b83cb7f0011cdfc352d30a4 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Tue, 3 Feb 2026 10:49:49 +0530 Subject: [PATCH] fix: refactor, race condition handle, fstring sql injection --- .../key_management_endpoints.py | 19 ++++-- litellm/proxy/utils.py | 65 ++++--------------- 2 files changed, 28 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index dea1e55f7f9..55f96e4189c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3365,12 +3365,21 @@ async def regenerate_key_fn( # noqa: PLR0915 revoke_at = datetime.now(timezone.utc) + timedelta( hours=grace_period_hours ) - await prisma_client.db.litellm_deprecatedverificationtoken.create( + # Use upsert to handle concurrent rotations gracefully; avoids + # unique constraint violation if same key is rotated simultaneously + await prisma_client.db.litellm_deprecatedverificationtoken.upsert( + where={"token": hashed_api_key}, data={ - "token": hashed_api_key, - "active_token_id": new_token_hash, - "revoke_at": revoke_at, - } + "create": { + "token": hashed_api_key, + "active_token_id": new_token_hash, + "revoke_at": revoke_at, + }, + "update": { + "active_token_id": new_token_hash, + "revoke_at": revoke_at, + }, + }, ) verbose_proxy_logger.debug( "Deprecated key retained for %s hours (revoke_at: %s)", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c9b07ce447d..5cc5507dd09 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7,7 +7,7 @@ import smtplib import threading import time import traceback -from datetime import date, datetime, timedelta +from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText from typing import ( @@ -2647,57 +2647,20 @@ class PrismaClient: # If not found in main table, check deprecated keys (grace period) if response is None: try: - deprecated_sql = f""" - SELECT active_token_id FROM "LiteLLM_DeprecatedVerificationToken" - WHERE token = '{hashed_token}' AND revoke_at > NOW() - LIMIT 1 - """ - deprecated_row = ( - await self._query_first_with_cached_plan_fallback( - deprecated_sql - ) + deprecated_row = await self.db.litellm_deprecatedverificationtoken.find_first( + where={ + "token": hashed_token, + "revoke_at": {"gt": datetime.now(timezone.utc)}, + }, + select={"active_token_id": True}, ) - if deprecated_row and deprecated_row.get("active_token_id"): - active_token_id = deprecated_row["active_token_id"] - sql_query_active = f""" - SELECT - v.*, - t.spend AS team_spend, - t.max_budget AS team_max_budget, - t.tpm_limit AS team_tpm_limit, - t.rpm_limit AS team_rpm_limit, - t.models AS team_models, - t.metadata AS team_metadata, - t.blocked AS team_blocked, - t.team_alias AS team_alias, - t.metadata AS team_metadata, - t.members_with_roles AS team_members_with_roles, - t.object_permission_id AS team_object_permission_id, - t.organization_id as org_id, - tm.spend AS team_member_spend, - m.aliases AS team_model_aliases, - b.max_budget AS litellm_budget_table_max_budget, - b.tpm_limit AS litellm_budget_table_tpm_limit, - b.rpm_limit AS litellm_budget_table_rpm_limit, - b.model_max_budget as litellm_budget_table_model_max_budget, - b.soft_budget as litellm_budget_table_soft_budget, - o.metadata as organization_metadata, - b2.max_budget as organization_max_budget, - b2.tpm_limit as organization_tpm_limit, - b2.rpm_limit as organization_rpm_limit - FROM "LiteLLM_VerificationToken" AS v - LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id - LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id - LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id - LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id - LEFT JOIN "LiteLLM_OrganizationTable" AS o ON v.organization_id = o.organization_id - LEFT JOIN "LiteLLM_BudgetTable" AS b2 ON o.budget_id = b2.budget_id - WHERE v.token = '{active_token_id}' - """ - response = ( - await self._query_first_with_cached_plan_fallback( - sql_query_active - ) + if deprecated_row and deprecated_row.active_token_id: + response = await self.get_data( + token=deprecated_row.active_token_id, + table_name="combined_view", + query_type="find_unique", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) if response is not None: verbose_proxy_logger.debug(