scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

fix: refactor, race condition handle, fstring sql injection

Browse source
This commit is contained in:
Harshit Jain 2026-02-03 10:49:49 +05:30
parent 768f9a44b2
commit 65401138b5
No known key found for this signature in database
GPG key ID: 36C392CD4415B4CF
2 changed files with 28 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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)",

65
litellm/proxy/utils.py
View file

@ -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(