mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
9b7ebb6a7d
* build(pyproject.toml): add new dev dependencies - for type checking * build: reformat files to fit black * ci: reformat to fit black * ci(test-litellm.yml): make tests run clear * build(pyproject.toml): add ruff * fix: fix ruff checks * build(mypy/): fix mypy linting errors * fix(hashicorp_secret_manager.py): fix passing cert for tls auth * build(mypy/): resolve all mypy errors * test: update test * fix: fix black formatting * build(pre-commit-config.yaml): use poetry run black * fix(proxy_server.py): fix linting error * fix: fix ruff safe representation error
106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
# 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._logging import verbose_proxy_logger
|
|
from litellm.secret_managers.aws_secret_manager import decrypt_env_var
|
|
|
|
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!")
|