(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
This commit is contained in:
Ishaan Jaff 2025-02-20 18:30:23 -08:00 committed by GitHub
parent 300d7825f5
commit 55b938dd6e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 168 additions and 22 deletions

View file

@ -1112,6 +1112,23 @@ jobs:
working_directory: ~/project working_directory: ~/project
steps: steps:
- checkout - 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: - run:
name: Build Docker image name: Build Docker image
command: | command: |
@ -1119,29 +1136,48 @@ jobs:
- run: - run:
name: Run Docker container name: Run Docker container
command: | command: |
docker run --name my-app \ docker run -d \
-p 4000:4000 \ -p 4000:4000 \
-e DATABASE_URL=$PROXY_DATABASE_URL \ -e DATABASE_URL=$PROXY_DATABASE_URL \
-e DISABLE_SCHEMA_UPDATE="True" \ -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/schema.prisma \
-v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/litellm/proxy/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 \ -v $(pwd)/litellm/proxy/example_config_yaml/disable_schema_update.yaml:/app/config.yaml \
--name my-app \
myapp:latest \ myapp:latest \
--config /app/config.yaml \ --config /app/config.yaml \
--port 4000 > docker_output.log 2>&1 || true --port 4000
- run: - run:
name: Display Docker logs name: Install curl and dockerize
command: cat docker_output.log
- run:
name: Check for expected error
command: | command: |
if grep -q "prisma schema out of sync with db. Consider running these sql_commands to sync the two" docker_output.log; then sudo apt-get update
echo "Expected error found. Test passed." 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 else
echo "Expected error not found. Test failed." echo "Expected message not found in logs. Test failed."
cat docker_output.log
exit 1 exit 1
fi 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: build_and_test:
machine: machine:

View file

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

View file

@ -7,7 +7,7 @@ import os
import urllib import urllib
import urllib.parse import urllib.parse
from datetime import datetime, timedelta 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 from litellm.secret_managers.main import str_to_bool
@ -112,12 +112,28 @@ class PrismaWrapper:
return original_attr 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: if disable_updates is None:
disable_prisma_schema_update = str_to_bool(os.getenv("DISABLE_SCHEMA_UPDATE")) disable_updates = os.getenv("DISABLE_SCHEMA_UPDATE", "false")
if disable_prisma_schema_update is True:
return False if isinstance(disable_updates, str):
return True disable_updates = str_to_bool(disable_updates)
return not bool(disable_updates)

View file

@ -4,6 +4,11 @@ model_list:
model: openai/fake model: openai/fake
api_key: fake-key api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/ 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: litellm_settings:
callbacks: ["gcs_bucket"] callbacks: ["gcs_bucket"]

View file

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

View file

@ -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

View file

@ -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