(Infra/DB) - Allow running older litellm version when out of sync with current state of DB (#8695)

* fix check migration

* clean up should_update_prisma_schema

* update test

* db_migration_disable_update_check

* Check container logs for expected message

* db_migration_disable_update_check

* test_check_migration_out_of_sync

* test_should_update_prisma_schema

* db_migration_disable_update_check

* pip install aiohttp

.circleci/config.yml

@@ -1112,6 +1112,23 @@ jobs:
     working_directory: ~/project
     steps:
       - checkout
+      - run:
+          name: Install Python 3.9
+          command: |
+            curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh
+            bash miniconda.sh -b -p $HOME/miniconda
+            export PATH="$HOME/miniconda/bin:$PATH"
+            conda init bash
+            source ~/.bashrc
+            conda create -n myenv python=3.9 -y
+            conda activate myenv
+            python --version
+      - run:
+          name: Install Dependencies
+          command: |
+            pip install "pytest==7.3.1"
+            pip install "pytest-asyncio==0.21.1"
+            pip install aiohttp
       - run:
           name: Build Docker image
           command: |
@@ -1119,29 +1136,48 @@ jobs:
       - run:
           name: Run Docker container
           command: |
-            docker run --name my-app \
+            docker run -d \
               -p 4000:4000 \
               -e DATABASE_URL=$PROXY_DATABASE_URL \
               -e DISABLE_SCHEMA_UPDATE="True" \
               -v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/schema.prisma \
               -v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/litellm/proxy/schema.prisma \
               -v $(pwd)/litellm/proxy/example_config_yaml/disable_schema_update.yaml:/app/config.yaml \
+              --name my-app \
               myapp:latest \
               --config /app/config.yaml \
-              --port 4000 > docker_output.log 2>&1 || true
+              --port 4000
       - run:
-          name: Display Docker logs
-          command: cat docker_output.log
-      - run:
-          name: Check for expected error
+          name: Install curl and dockerize
           command: |
-            if grep -q "prisma schema out of sync with db. Consider running these sql_commands to sync the two" docker_output.log; then
-              echo "Expected error found. Test passed."
+            sudo apt-get update
+            sudo apt-get install -y curl
+            sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz
+            sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz
+            sudo rm dockerize-linux-amd64-v0.6.1.tar.gz
+
+      - run:
+          name: Wait for container to be ready
+          command: dockerize -wait http://localhost:4000 -timeout 1m
+      - run:
+          name: Check container logs for expected message
+          command: |
+            echo "=== Printing Full Container Startup Logs ==="
+            docker logs my-app
+            echo "=== End of Full Container Startup Logs ==="
+            
+            if docker logs my-app 2>&1 | grep -q "prisma schema out of sync with db. Consider running these sql_commands to sync the two"; then
+              echo "Expected message found in logs. Test passed."
             else
-              echo "Expected error not found. Test failed."
-              cat docker_output.log
+              echo "Expected message not found in logs. Test failed."
               exit 1
             fi
+      - run:
+          name: Run Basic Proxy Startup Tests (Health Readiness and Chat Completion)
+          command: |
+            python -m pytest -vv tests/basic_proxy_startup_tests -x --junitxml=test-results/junit-2.xml --durations=5
+          no_output_timeout: 120m
+
 
   build_and_test:
     machine:

litellm/proxy/db/check_migration.py

@@ -4,6 +4,8 @@ import os
 import subprocess
 from typing import List, Optional, Tuple
 
+from litellm._logging import verbose_logger
+
 
 def extract_sql_commands(diff_output: str) -> List[str]:
     """
@@ -52,6 +54,7 @@ def check_prisma_schema_diff_helper(db_url: str) -> Tuple[bool, List[str]]:
         subprocess.CalledProcessError: If the Prisma command fails.
         Exception: For any other errors during execution.
     """
+    verbose_logger.debug("Checking for Prisma schema diff...")  # noqa: T201
     try:
         result = subprocess.run(
             [
@@ -94,8 +97,8 @@ def check_prisma_schema_diff(db_url: Optional[str] = None) -> None:
             raise Exception("DATABASE_URL not set")
     has_diff, message = check_prisma_schema_diff_helper(db_url)
     if has_diff:
-        raise Exception(
-            "prisma schema out of sync with db. Consider running these sql_commands to sync the two - {}".format(
+        verbose_logger.exception(
+            "🚨🚨🚨 prisma schema out of sync with db. Consider running these sql_commands to sync the two - {}".format(
                 message
             )
         )

litellm/proxy/db/prisma_client.py

@@ -7,7 +7,7 @@ import os
 import urllib
 import urllib.parse
 from datetime import datetime, timedelta
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from litellm.secret_managers.main import str_to_bool
 
@@ -112,12 +112,28 @@ class PrismaWrapper:
         return original_attr
 
 
-def should_update_schema(disable_prisma_schema_update: Optional[bool]):
+def should_update_prisma_schema(
+    disable_updates: Optional[Union[bool, str]] = None
+) -> bool:
     """
-    This function is used to determine if the Prisma schema should be updated.
+    Determines if Prisma Schema updates should be applied during startup.
+
+    Args:
+        disable_updates: Controls whether schema updates are disabled.
+            Accepts boolean or string ('true'/'false'). Defaults to checking DISABLE_SCHEMA_UPDATE env var.
+
+    Returns:
+        bool: True if schema updates should be applied, False if updates are disabled.
+
+    Examples:
+        >>> should_update_prisma_schema()  # Checks DISABLE_SCHEMA_UPDATE env var
+        >>> should_update_prisma_schema(True)  # Explicitly disable updates
+        >>> should_update_prisma_schema("false")  # Enable updates using string
     """
-    if disable_prisma_schema_update is None:
-        disable_prisma_schema_update = str_to_bool(os.getenv("DISABLE_SCHEMA_UPDATE"))
-    if disable_prisma_schema_update is True:
-        return False
-    return True
+    if disable_updates is None:
+        disable_updates = os.getenv("DISABLE_SCHEMA_UPDATE", "false")
+
+    if isinstance(disable_updates, str):
+        disable_updates = str_to_bool(disable_updates)
+
+    return not bool(disable_updates)

litellm/proxy/example_config_yaml/disable_schema_update.yaml

@@ -4,6 +4,11 @@ model_list:
       model: openai/fake
       api_key: fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
+  - model_name: gpt-4
+    litellm_params:
+      model: openai/gpt-4
+      api_key: fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
 
 litellm_settings:
   callbacks: ["gcs_bucket"]

litellm/proxy/proxy_cli.py

@@ -607,10 +607,10 @@ def run_server(  # noqa: PLR0915
 
             if is_prisma_runnable:
                 from litellm.proxy.db.check_migration import check_prisma_schema_diff
-                from litellm.proxy.db.prisma_client import should_update_schema
+                from litellm.proxy.db.prisma_client import should_update_prisma_schema
 
                 if (
-                    should_update_schema(
+                    should_update_prisma_schema(
                         general_settings.get("disable_prisma_schema_update")
                     )
                     is False

tests/litellm/proxy/db/test_check_migration.py

@@ -0,0 +1,51 @@
+import json
+import os
+import sys
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../../..")
+)  # Adds the parent directory to the system path
+
+
+import json
+import os
+import sys
+import time
+
+import pytest
+from fastapi.testclient import TestClient
+
+import litellm
+
+
+def test_check_migration_out_of_sync(mocker):
+    """
+    Test that the check_prisma_schema_diff function
+    - 🚨 [IMPORTANT] Does NOT Raise an Exception when the Prisma schema is out of sync with the database.
+    - logs an error when the Prisma schema is out of sync with the database.
+    """
+    # Mock the logger BEFORE importing the function
+    mock_logger = mocker.patch("litellm._logging.verbose_logger")
+
+    # Import the function after mocking the logger
+    from litellm.proxy.db.check_migration import check_prisma_schema_diff
+
+    # Mock the helper function to simulate out-of-sync state
+    mock_diff_helper = mocker.patch(
+        "litellm.proxy.db.check_migration.check_prisma_schema_diff_helper",
+        return_value=(True, ["ALTER TABLE users ADD COLUMN new_field TEXT;"]),
+    )
+
+    # Run the function - it should not raise an error
+    try:
+        check_prisma_schema_diff(db_url="mock_url")
+    except Exception as e:
+        pytest.fail(f"check_prisma_schema_diff raised an unexpected exception: {e}")
+
+    # Verify the logger was called with the expected message
+    mock_logger.exception.assert_called_once()
+    actual_message = mock_logger.exception.call_args[0][0]
+    assert "prisma schema out of sync with db" in actual_message

tests/litellm/proxy/db/test_prisma_client.py

@@ -0,0 +1,35 @@
+import json
+import os
+import sys
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../../..")
+)  # Adds the parent directory to the system path
+
+
+from litellm.proxy.db.prisma_client import should_update_prisma_schema
+
+
+def test_should_update_prisma_schema(monkeypatch):
+    # CASE 1: Environment variable behavior
+    # When DISABLE_SCHEMA_UPDATE is not set -> should update
+    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
+    assert should_update_prisma_schema() == True
+
+    # When DISABLE_SCHEMA_UPDATE="true" -> should not update
+    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "true")
+    assert should_update_prisma_schema() == False
+
+    # When DISABLE_SCHEMA_UPDATE="false" -> should update
+    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "false")
+    assert should_update_prisma_schema() == True
+
+    # CASE 2: Explicit parameter behavior (overrides env var)
+    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
+    assert should_update_prisma_schema(True) == False  # Param True -> should not update
+
+    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)  # Set env var opposite to param
+    assert should_update_prisma_schema(False) == True  # Param False -> should update