forked from BerriAI/litellm
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathupdate_unassigned_teams.py
34 lines (32 loc) · 1.01 KB
/
update_unassigned_teams.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from prisma import Prisma
from litellm._logging import verbose_logger
async def apply_db_fixes(db: Prisma):
"""
Do Not Run this in production, only use it as a one-time fix
"""
verbose_logger.warning(
"DO NOT run this in Production....Running update_unassigned_teams"
)
try:
sql_query = """
UPDATE "LiteLLM_SpendLogs"
SET team_id = (
SELECT vt.team_id
FROM "LiteLLM_VerificationToken" vt
WHERE vt.token = "LiteLLM_SpendLogs".api_key
)
WHERE team_id IS NULL
AND EXISTS (
SELECT 1
FROM "LiteLLM_VerificationToken" vt
WHERE vt.token = "LiteLLM_SpendLogs".api_key
);
"""
response = await db.query_raw(sql_query)
print(
"Updated unassigned teams, Response=%s",
response,
)
except Exception as e:
raise Exception(f"Error apply_db_fixes: {str(e)}")
return