litellm-proxy-extras 0.1.12__tar.gz → 0.1.13__tar.gz

{litellm_proxy_extras-0.1.12 → litellm_proxy_extras-0.1.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: litellm-proxy-extras
-Version: 0.1.12
+Version: 0.1.13
 Summary: Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package.
 Author: BerriAI
 Requires-Python: >=3.8, !=2.7.*, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*, !=3.7.*

{litellm_proxy_extras-0.1.12 → litellm_proxy_extras-0.1.13}/litellm_proxy_extras/utils.py

@@ -2,8 +2,10 @@ import glob
 import os
 import random
 import re
+import shutil
 import subprocess
 import time
+from datetime import datetime
 from pathlib import Path
 from typing import Optional
 
@@ -19,9 +21,30 @@ def str_to_bool(value: Optional[str]) -> bool:
 class ProxyExtrasDBManager:
     @staticmethod
     def _get_prisma_dir() -> str:
-        """Get the path to the migrations directory"""
-        migrations_dir = os.path.dirname(__file__)
-        return migrations_dir
+        """
+        Get the path to the migrations directory
+
+        Set os.environ["LITELLM_MIGRATION_DIR"] to a custom migrations directory, to support baselining db in read-only fs.
+        """
+        custom_migrations_dir = os.getenv("LITELLM_MIGRATION_DIR")
+        pkg_migrations_dir = os.path.dirname(__file__)
+        if custom_migrations_dir:
+            # If migrations_dir exists, copy contents
+            if os.path.exists(custom_migrations_dir):
+                # Copy contents instead of directory itself
+                for item in os.listdir(pkg_migrations_dir):
+                    src_path = os.path.join(pkg_migrations_dir, item)
+                    dst_path = os.path.join(custom_migrations_dir, item)
+                    if os.path.isdir(src_path):
+                        shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
+                    else:
+                        shutil.copy2(src_path, dst_path)
+            else:
+                # If directory doesn't exist, create it and copy everything
+                shutil.copytree(pkg_migrations_dir, custom_migrations_dir)
+            return custom_migrations_dir
+
+        return pkg_migrations_dir
 
     @staticmethod
     def _create_baseline_migration(schema_path: str) -> bool:
@@ -33,27 +56,29 @@ class ProxyExtrasDBManager:
         # Create migrations/0_init directory
         init_dir.mkdir(parents=True, exist_ok=True)
 
-        # Generate migration SQL file
-        migration_file = init_dir / "migration.sql"
+        database_url = os.getenv("DATABASE_URL")
 
         try:
-            # Generate migration diff with increased timeout
+            # 1. Generate migration SQL file by comparing empty state to current db state
+            logger.info("Generating baseline migration...")
+            migration_file = init_dir / "migration.sql"
             subprocess.run(
                 [
                     "prisma",
                     "migrate",
                     "diff",
                     "--from-empty",
-                    "--to-schema-datamodel",
-                    str(schema_path),
+                    "--to-url",
+                    database_url,
                     "--script",
                 ],
                 stdout=open(migration_file, "w"),
                 check=True,
                 timeout=30,
-            )  # 30 second timeout
+            )
 
-            # Mark migration as applied with increased timeout
+            # 3. Mark the migration as applied since it represents current state
+            logger.info("Marking baseline migration as applied...")
             subprocess.run(
                 [
                     "prisma",
@@ -73,8 +98,10 @@ class ProxyExtrasDBManager:
             )
             return False
         except subprocess.CalledProcessError as e:
-            logger.warning(f"Error creating baseline migration: {e}")
-            return False
+            logger.warning(
+                f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}"
+            )
+            raise e
 
     @staticmethod
     def _get_migration_names(migrations_dir: str) -> list:
@@ -104,8 +131,85 @@ class ProxyExtrasDBManager:
         )
 
     @staticmethod
-    def _resolve_all_migrations(migrations_dir: str):
-        """Mark all existing migrations as applied"""
+    def _resolve_all_migrations(migrations_dir: str, schema_path: str):
+        """
+        1. Compare the current database state to schema.prisma and generate a migration for the diff.
+        2. Run prisma migrate deploy to apply any pending migrations.
+        3. Mark all existing migrations as applied.
+        """
+        database_url = os.getenv("DATABASE_URL")
+        diff_dir = (
+            Path(migrations_dir)
+            / "migrations"
+            / f"{datetime.now().strftime('%Y%m%d%H%M%S')}_baseline_diff"
+        )
+        try:
+            diff_dir.mkdir(parents=True, exist_ok=True)
+        except Exception as e:
+            if "Permission denied" in str(e):
+                logger.warning(
+                    f"Permission denied - {e}\nunable to baseline db. Set LITELLM_MIGRATION_DIR environment variable to a writable directory to enable migrations."
+                )
+                return
+            raise e
+        diff_sql_path = diff_dir / "migration.sql"
+
+        # 1. Generate migration SQL for the diff between DB and schema
+        try:
+            logger.info("Generating migration diff between DB and schema.prisma...")
+            with open(diff_sql_path, "w") as f:
+                subprocess.run(
+                    [
+                        "prisma",
+                        "migrate",
+                        "diff",
+                        "--from-url",
+                        database_url,
+                        "--to-schema-datamodel",
+                        schema_path,
+                        "--script",
+                    ],
+                    check=True,
+                    timeout=60,
+                    stdout=f,
+                )
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to generate migration diff: {e.stderr}")
+        except subprocess.TimeoutExpired:
+            logger.warning("Migration diff generation timed out.")
+
+        # check if the migration was created
+        if not diff_sql_path.exists():
+            logger.warning("Migration diff was not created")
+            return
+        logger.info(f"Migration diff created at {diff_sql_path}")
+
+        # 2. Run prisma db execute to apply the migration
+        try:
+            logger.info("Running prisma db execute to apply the migration diff...")
+            result = subprocess.run(
+                [
+                    "prisma",
+                    "db",
+                    "execute",
+                    "--file",
+                    str(diff_sql_path),
+                    "--schema",
+                    schema_path,
+                ],
+                timeout=60,
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            logger.info(f"prisma db execute stdout: {result.stdout}")
+            logger.info("✅ Migration diff applied successfully")
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to apply migration diff: {e.stderr}")
+        except subprocess.TimeoutExpired:
+            logger.warning("Migration diff application timed out.")
+
+        # 3. Mark all migrations as applied
         migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
         logger.info(f"Resolving {len(migration_names)} migrations")
         for migration_name in migration_names:
@@ -126,7 +230,7 @@ class ProxyExtrasDBManager:
                     )
 
     @staticmethod
-    def setup_database(schema_path: str, use_migrate: bool = False) -> bool:
+    def setup_database(use_migrate: bool = False) -> bool:
         """
         Set up the database using either prisma migrate or prisma db push
         Uses migrations from litellm-proxy-extras package
@@ -138,6 +242,7 @@ class ProxyExtrasDBManager:
         Returns:
             bool: True if setup was successful, False otherwise
         """
+        schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
         use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate
         for attempt in range(4):
             original_dir = os.getcwd()
@@ -200,7 +305,9 @@ class ProxyExtrasDBManager:
                             logger.info(
                                 "Baseline migration created, resolving all migrations"
                             )
-                            ProxyExtrasDBManager._resolve_all_migrations(migrations_dir)
+                            ProxyExtrasDBManager._resolve_all_migrations(
+                                migrations_dir, schema_path
+                            )
                             logger.info("✅ All migrations resolved.")
                             return True
                         elif (

{litellm_proxy_extras-0.1.12 → litellm_proxy_extras-0.1.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm-proxy-extras"
-version = "0.1.12"
+version = "0.1.13"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 authors = ["BerriAI"]
 readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "0.1.12"
+version = "0.1.13"
 version_files = [
     "pyproject.toml:version",
     "../requirements.txt:litellm-proxy-extras==",