diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index 6df123bfe8f..6e2a06c96e1 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -47,19 +47,26 @@ def extract_sql_commands(diff_output: str) -> list[str]: def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: """Checks for differences between current database and Prisma schema. - Never raises: a diff that cannot be produced, because the command failed or - because it outlived its budget, is reported as "no diff" so boot continues. + Never raises: a diff that cannot be produced, because the runner is missing, + because the command failed, or because it outlived its budget, is reported as + "no diff" so boot continues. Returns: A tuple containing: - A boolean indicating if differences were found (True) or not (False). - The SQL commands that would close the diff, empty when there is none. """ - from litellm_proxy_extras.prisma_toolchain import ( - PRISMA_COMMAND_TIMEOUT_ENV_VAR, - prisma_command_timeout, - run_prisma, - ) + try: + from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" + ) + return False, [] verbose_logger.debug("Checking for Prisma schema diff...") timeout: Final = prisma_command_timeout() diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 82cf910d2f5..18663e6317e 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -937,10 +937,14 @@ class PrismaManager: use_v2_resolver=use_v2_resolver, ) else: - from litellm_proxy_extras.prisma_toolchain import ( - prisma_command_timeout, - run_prisma, - ) + try: + from litellm_proxy_extras.prisma_toolchain import ( + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e) + return False PrismaManager._raise_if_partitioned_spend_logs() run_prisma( diff --git a/tests/test_litellm/proxy/db/test_check_migration.py b/tests/test_litellm/proxy/db/test_check_migration.py index 36c970de5ad..74a6841cb23 100644 --- a/tests/test_litellm/proxy/db/test_check_migration.py +++ b/tests/test_litellm/proxy/db/test_check_migration.py @@ -55,3 +55,17 @@ def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fak "--to-schema-datamodel", "./schema.prisma", "--script"] ] assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch): + """ + Boot calls this helper directly, so an ImportError here takes the proxy down before + uvicorn starts. An install without the runner must lose the diagnostic, not the proxy. + """ + import sys + + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index fe348858950..7f43e483557 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -387,3 +387,16 @@ def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_d assert PrismaManager.setup_database(use_migrate=False) is True assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV] assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot( + fake_prisma_cli, unset_database_url, monkeypatch +): + """ + An ImportError out of setup_database escapes the caller's RuntimeError handler and + kills boot, bypassing the operator's enforce_prisma_migration_check choice. + """ + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert PrismaManager.setup_database(use_migrate=False) is False + assert fake_prisma_cli.calls == []