Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,2 +1 @@
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");
SELECT 1;
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Original file line number Diff line number Diff line change
@@ -1,12 +1 @@
-- CreateIndex (CONCURRENTLY)
--
-- Disclaimer:
-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
-- interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
-- - Do not edit this file after it has been applied to any database: Prisma checksums
-- migrations; add a new migration instead.
-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
-- without IF NOT EXISTS if you must support older versions).
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
SELECT 1;
Comment thread
greptile-apps[bot] marked this conversation as resolved.
13 changes: 11 additions & 2 deletions litellm/proxy/prisma_migration.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,24 @@

from typing import Final

import click
from litellm_proxy_extras.prisma_toolchain import resolve_prisma_argv

from litellm._logging import verbose_proxy_logger
from litellm.proxy.proxy_cli import run_server
from litellm.secret_managers.main import str_to_bool


def _enforce_prisma_migration_check(raw: str | None) -> bool:
if raw is None or not raw.strip():
return True
try:
return bool(click.BOOL(raw))
except click.BadParameter:
return True


def main() -> int:
enforce_prisma_migration_check: Final = str_to_bool(os.getenv("ENFORCE_PRISMA_MIGRATION_CHECK")) is not False
enforce_prisma_migration_check: Final = _enforce_prisma_migration_check(os.getenv("ENFORCE_PRISMA_MIGRATION_CHECK"))
run_server_args: Final = (
("--skip_server_startup", "--enforce_prisma_migration_check")
if enforce_prisma_migration_check
Expand Down
16 changes: 10 additions & 6 deletions litellm/proxy/proxy_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -963,10 +963,14 @@ def _maybe_setup_prometheus_multiproc_dir(
envvar="LIMIT_CONCURRENCY",
)
@click.option(
"--enforce_prisma_migration_check",
is_flag=True,
default=False,
help="Exit with error if database migration fails on startup.",
"--enforce_prisma_migration_check/--no-enforce_prisma_migration_check",
default=True,
Comment thread
greptile-apps[bot] marked this conversation as resolved.
show_default=True,
help=(
"Exit when database setup fails on startup instead of serving against a database "
"whose schema may be behind the code. Opt out with --no-enforce_prisma_migration_check "
"or ENFORCE_PRISMA_MIGRATION_CHECK=false."
),
envvar="ENFORCE_PRISMA_MIGRATION_CHECK",
)
@click.option(
Expand Down Expand Up @@ -1434,8 +1438,8 @@ def run_server(
sys.exit(1)
else:
print(
"\033[1;33mLiteLLM Proxy: Database migration failed but continuing startup. "
"Set --enforce_prisma_migration_check or ENFORCE_PRISMA_MIGRATION_CHECK=true to exit on failure.\033[0m"
"\033[1;33mLiteLLM Proxy: Database migration failed but continuing startup because "
"ENFORCE_PRISMA_MIGRATION_CHECK is disabled. The schema may be behind the code.\033[0m"
)
else:
print(
Expand Down
34 changes: 32 additions & 2 deletions tests/test_litellm/proxy/test_prisma_migration.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,32 @@ def test_main_disables_migration_check_when_explicitly_false(

mock_run_server.assert_called_once_with(("--skip_server_startup",), standalone_mode=False)

@pytest.mark.parametrize(
("value", "enforces"),
[
("0", False),
("no", False),
("off", False),
("FALSE", False),
("1", True),
("yes", True),
("", True),
("maybe", True),
],
)
@patch("litellm.proxy.prisma_migration.subprocess.run")
@patch("litellm.proxy.prisma_migration.run_server")
def test_main_reads_the_same_opt_out_values_as_the_proxy_flag(
self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, value: str, enforces: bool
) -> None:
mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="")

with patch.dict(os.environ, {"ENFORCE_PRISMA_MIGRATION_CHECK": value}, clear=True):
assert prisma_migration.main() == 0

passed: tuple[str, ...] = mock_run_server.call_args.args[0]
assert ("--enforce_prisma_migration_check" in passed) is enforces

@pytest.mark.parametrize("env", [{}, {"ENFORCE_PRISMA_MIGRATION_CHECK": "false"}])
@patch("litellm.proxy.prisma_migration.subprocess.run")
@patch("litellm.proxy.prisma_migration.run_server")
Expand Down Expand Up @@ -67,7 +93,9 @@ def test_main_propagates_migration_failure(

mock_subprocess_run.assert_not_called()

@patch("litellm.proxy.prisma_migration.subprocess.run") # test-quality-ok: the spawned argv is the behavior under test
@patch(
"litellm.proxy.prisma_migration.subprocess.run"
) # test-quality-ok: the spawned argv is the behavior under test
@patch("litellm.proxy.prisma_migration.run_server") # test-quality-ok: run_server boots the whole proxy
def test_prisma_generate_runs_through_the_module_when_the_cli_is_not_on_path(
self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path
Expand All @@ -81,7 +109,9 @@ def test_prisma_generate_runs_through_the_module_when_the_cli_is_not_on_path(

assert mock_subprocess_run.call_args.args[0] == (sys.executable, "-m", "prisma", "generate")

@patch("litellm.proxy.prisma_migration.subprocess.run") # test-quality-ok: the spawned argv is the behavior under test
@patch(
"litellm.proxy.prisma_migration.subprocess.run"
) # test-quality-ok: the spawned argv is the behavior under test
@patch("litellm.proxy.prisma_migration.run_server") # test-quality-ok: run_server boots the whole proxy
def test_prisma_generate_runs_the_console_script_when_it_is_on_path(
self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path
Expand Down
43 changes: 27 additions & 16 deletions tests/test_litellm/proxy/test_proxy_cli.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import inspect
import os
from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
Expand Down Expand Up @@ -2266,6 +2267,18 @@ def test_migrations_run_when_the_prisma_cli_is_not_on_path(
use_migrate=True, use_v2_resolver=True
)

@pytest.mark.parametrize(
("arguments", "env", "exits"),
[
((), {}, True),
(("--enforce_prisma_migration_check",), {}, True),
((), {"ENFORCE_PRISMA_MIGRATION_CHECK": "true"}, True),
(("--no-enforce_prisma_migration_check",), {}, False),
((), {"ENFORCE_PRISMA_MIGRATION_CHECK": "false"}, False),
((), {"ENFORCE_PRISMA_MIGRATION_CHECK": "0"}, False),
],
ids=("default", "flag", "env-true", "no-flag", "env-false", "env-0"),
)
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
Expand All @@ -2278,8 +2291,13 @@ def test_startup_fails_when_db_setup_fails(
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
arguments,
env,
exits,
capsys,
):
"""Test that proxy exits with code 1 when PrismaManager.setup_database returns False and --enforce_prisma_migration_check is set"""
"""A failed PrismaManager.setup_database exits 1 by default; only an explicit
--no-enforce_prisma_migration_check or ENFORCE_PRISMA_MIGRATION_CHECK=false keeps the proxy booting"""
from litellm.proxy.proxy_cli import run_server

mock_subprocess_run.return_value = MagicMock(returncode=0)
Expand All @@ -2296,12 +2314,13 @@ def test_startup_fails_when_db_setup_fails(
clean_env = {
k: v
for k, v in os.environ.items()
if k not in ("DATABASE_URL", "DIRECT_URL")
if k not in ("DATABASE_URL", "DIRECT_URL", "ENFORCE_PRISMA_MIGRATION_CHECK")
}
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
outcome = pytest.raises(SystemExit) if exits else nullcontext()

with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(os.environ, {**clean_env, **env}, clear=True),
patch.dict(
"sys.modules",
{
Expand All @@ -2312,26 +2331,18 @@ def test_startup_fails_when_db_setup_fails(
patch(
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
) as mock_get_args,
outcome as exc_info,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
run_server.main(["--local", "--skip_server_startup", *arguments], standalone_mode=False)

with pytest.raises(SystemExit) as exc_info:
run_server.main(
[
"--local",
"--skip_server_startup",
"--enforce_prisma_migration_check",
],
standalone_mode=False,
)
assert exc_info.value.code == 1
mock_setup_database.assert_called_once_with(
use_migrate=True, use_v2_resolver=True
)
assert (exc_info is not None and exc_info.value.code == 1) is exits
assert ("Database migration failed but continuing startup" in capsys.readouterr().out) is not exits
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)

@patch("subprocess.run")
@patch("atexit.register")
Expand Down
Loading