From 58d988210233b34807188ef5526b2060874b4e4c Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 14 Jul 2025 22:11:46 -0700
Subject: [PATCH] refactor(prisma_migration.py): refactor to support
 use_prisma_migrate - for helm hook (#12600)

* refactor(prisma_migration.py): refactor to support use_prisma_migrate for helm hook

* fix(prisma_migration.py): don't use subprocess

* fix(prisma_migration.py): fix cli commands

* fix(prisma_migration.py): still run prisma generate
---
 litellm/proxy/prisma_migration.py | 104 ++++--------------------------
 1 file changed, 13 insertions(+), 91 deletions(-)

diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py
index c1e2220c159..35142736ab3 100644
--- a/litellm/proxy/prisma_migration.py
+++ b/litellm/proxy/prisma_migration.py
@@ -4,103 +4,25 @@
 import os
 import subprocess
 import sys
-import time
 
 sys.path.insert(
     0, os.path.abspath("./")
 )  # Adds the parent directory to the system path
+
 from litellm._logging import verbose_proxy_logger
-from litellm.secret_managers.aws_secret_manager import decrypt_env_var
+from litellm.proxy.proxy_cli import run_server
 
-if os.getenv("USE_AWS_KMS", None) is not None and os.getenv("USE_AWS_KMS") == "True":
-    ## V2 IMPLEMENTATION OF AWS KMS - USER WANTS TO DECRYPT MULTIPLE KEYS IN THEIR ENV
-    new_env_var = decrypt_env_var()
+# Call the Click command with standalone_mode=False
+run_server(["--use_prisma_migrate", "--skip_server_startup"], standalone_mode=False)
 
-    for k, v in new_env_var.items():
-        os.environ[k] = v
+# run prisma generate
+verbose_proxy_logger.info("Running 'prisma generate'...")
+result = subprocess.run(["prisma", "generate"], capture_output=True, text=True)
+verbose_proxy_logger.info(f"'prisma generate' stdout: {result.stdout}")  # Log stdout
+exit_code = result.returncode
 
-# Check if DATABASE_URL is not set
-database_url = os.getenv("DATABASE_URL")
-if not database_url:
-    verbose_proxy_logger.info("Constructing DATABASE_URL from environment variables")
-    # Check if all required variables are provided
-    database_host = os.getenv("DATABASE_HOST")
-    database_username = os.getenv("DATABASE_USERNAME")
-    database_password = os.getenv("DATABASE_PASSWORD")
-    database_name = os.getenv("DATABASE_NAME")
-
-    if database_host and database_username and database_password and database_name:
-        # Construct DATABASE_URL from the provided variables
-        database_url = f"postgresql://{database_username}:{database_password}@{database_host}/{database_name}"
-        os.environ["DATABASE_URL"] = database_url  # Log the constructed URL
-    else:
-        verbose_proxy_logger.error(
-            "Error: Required database environment variables are not set. Provide a postgres url for DATABASE_URL."  # noqa
-        )
-        exit(1)
-else:
-    verbose_proxy_logger.info(
-        "Using existing DATABASE_URL environment variable"
-    )  # Log existing DATABASE_URL
-
-# Set DIRECT_URL to the value of DATABASE_URL if it is not set, required for migrations
-direct_url = os.getenv("DIRECT_URL")
-if not direct_url:
-    os.environ["DIRECT_URL"] = database_url
-
-# Apply migrations
-retry_count = 0
-max_retries = 3
-exit_code = 1
-
-disable_schema_update = os.getenv("DISABLE_SCHEMA_UPDATE")
-if disable_schema_update is not None and disable_schema_update == "True":
-    verbose_proxy_logger.info("Skipping schema update...")
-    exit(0)
-
-while retry_count < max_retries and exit_code != 0:
-    retry_count += 1
-    verbose_proxy_logger.info(f"Attempt {retry_count}...")
-
-    # run prisma generate
-    verbose_proxy_logger.info("Running 'prisma generate'...")
-    result = subprocess.run(["prisma", "generate"], capture_output=True, text=True)
-    verbose_proxy_logger.info(
-        f"'prisma generate' stdout: {result.stdout}"
-    )  # Log stdout
-    exit_code = result.returncode
-
-    if exit_code != 0:
-        verbose_proxy_logger.info(
-            f"'prisma generate' failed with exit code {exit_code}."
-        )
-        verbose_proxy_logger.error(
-            f"'prisma generate' stderr: {result.stderr}"
-        )  # Log stderr
-
-    # Run the Prisma db push command
-    verbose_proxy_logger.info("Running 'prisma db push --accept-data-loss'...")
-    result = subprocess.run(
-        ["prisma", "db", "push", "--accept-data-loss"], capture_output=True, text=True
-    )
-    verbose_proxy_logger.info(f"'prisma db push' stdout: {result.stdout}")  # Log stdout
-    exit_code = result.returncode
-
-    if exit_code != 0:
-        verbose_proxy_logger.info(
-            f"'prisma db push' stderr: {result.stderr}"
-        )  # Log stderr
-        verbose_proxy_logger.error(
-            f"'prisma db push' failed with exit code {exit_code}."
-        )
-        if retry_count < max_retries:
-            verbose_proxy_logger.info("Retrying in 10 seconds...")
-            time.sleep(10)
-
-if retry_count == max_retries and exit_code != 0:
+if exit_code != 0:
+    verbose_proxy_logger.info(f"'prisma generate' failed with exit code {exit_code}.")
     verbose_proxy_logger.error(
-        f"Unable to push database changes after {max_retries} retries."
-    )
-    exit(1)
-
-verbose_proxy_logger.info("Database push successful!")
+        f"'prisma generate' stderr: {result.stderr}"
+    )  # Log stderr