Spaces:
Sleeping
Sleeping
| # 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!") | |