1 files changed, 92 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/litellm/proxy/prisma_migration.py b/.venv/lib/python3.12/site-packages/litellm/proxy/prisma_migration.py
new file mode 100644
index 00000000..22fa4da9
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/litellm/proxy/prisma_migration.py
@@ -0,0 +1,92 @@
+# What is this?
+## Script to apply initial prisma migration on Docker setup
+
+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.secret_managers.aws_secret_manager import decrypt_env_var
+from litellm._logging import verbose_proxy_logger
+
+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()
+
+    for k, v in new_env_var.items():
+        os.environ[k] = v
+
+# 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:
+    verbose_proxy_logger.error(f"Unable to push database changes after {max_retries} retries.")
+    exit(1)
+
+verbose_proxy_logger.info("Database push successful!")