diff --git a/.circleci/config.yml b/.circleci/config.yml index cb887077de9..bade03f0058 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -853,6 +853,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: | @@ -860,30 +877,49 @@ 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: image: ubuntu-2204:2023.10.1 diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index ecb503db8ff..bf180c1132d 100644 --- a/litellm/proxy/db/check_migration.py +++ b/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 ) ) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 54d59bd3470..f8bb6b09fd9 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/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) diff --git a/litellm/proxy/example_config_yaml/disable_schema_update.yaml b/litellm/proxy/example_config_yaml/disable_schema_update.yaml index cc56b9516f2..5dcbd0dbd57 100644 --- a/litellm/proxy/example_config_yaml/disable_schema_update.yaml +++ b/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"] diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 5c4b04fb701..431e03ed15a 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -653,10 +653,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 diff --git a/tests/litellm/proxy/db/test_check_migration.py b/tests/litellm/proxy/db/test_check_migration.py new file mode 100644 index 00000000000..ad72a0d1195 --- /dev/null +++ b/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 diff --git a/tests/litellm/proxy/db/test_prisma_client.py b/tests/litellm/proxy/db/test_prisma_client.py new file mode 100644 index 00000000000..c7e99aa7540 --- /dev/null +++ b/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