From df1ac922223a2a6bab88d2292e27e55ec8c4f284 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 19 Jul 2024 12:49:23 -0700
Subject: [PATCH] fix add fix to update spend logs

---
 litellm/proxy/proxy_server.py |  3 +++
 litellm/proxy/utils.py        | 28 ++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index b3db5c720..79f25c6e1 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -2604,6 +2604,9 @@ async def startup_event():
     ### CHECK IF VIEW EXISTS ###
     if prisma_client is not None:
         create_view_response = await prisma_client.check_view_exists()
+        # Apply misc fixes on DB
+        # [non-blocking] helper to apply fixes from older litellm versions
+        asyncio.create_task(prisma_client.apply_db_fixes())
 
     ### START BATCH WRITING DB + CHECKING NEW MODELS###
     if prisma_client is not None:
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 8a8494ce4..f06b65705 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -1920,6 +1920,34 @@ class PrismaClient:
             )
             raise e
 
+    async def apply_db_fixes(self):
+        try:
+            verbose_proxy_logger.debug(
+                "Applying LiteLLM - DB Fixes fixing logs in SpendLogs"
+            )
+            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 self.db.query_raw(sql_query)
+            verbose_proxy_logger.debug(
+                "Applied LiteLLM - DB Fixes fixing logs in SpendLogs, Response=%s",
+                response,
+            )
+        except Exception as e:
+            verbose_proxy_logger.debug(f"Error apply_db_fixes: {str(e)}")
+        return
+
 
 class DBClient:
     """