diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index 6e2a06c96e1..c88bc1b4a1a 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -2,10 +2,41 @@ import os import subprocess +from collections.abc import Sequence from typing import Final from litellm._logging import verbose_logger +_SPEND_LOGS_TABLE: Final = '"LiteLLM_SpendLogs"' +_SPEND_LOGS_PKEY: Final = '"LiteLLM_SpendLogs_pkey"' + + +def _is_spend_logs_pkey_rewrite(command: str) -> bool: + """ + True when a suggested command only rewrites the LiteLLM_SpendLogs primary + key (DROP and/or ADD of LiteLLM_SpendLogs_pkey). + + schema.prisma declares request_id as a single-column @id, but with + use_spend_logs_partitioning the deployed table is partitioned by + "startTime" and its primary key must include the partition key - so the + startup migrate diff always proposes this rewrite, and running it on a + partitioned table is invalid. The quoted table name keeps + LiteLLM_SpendLogs_legacy and partition child tables out of the match. + """ + return _SPEND_LOGS_TABLE in command and _SPEND_LOGS_PKEY in command + + +def filter_spend_logs_pkey_rewrite( + sql_commands: Sequence[str], use_spend_logs_partitioning: bool +) -> list[str]: # mutable-ok: fresh filtered list for the caller + """ + Drop the SpendLogs primary-key rewrite from the suggested commands when + spend-log partitioning is on; every other detected drift still surfaces. + """ + if not use_spend_logs_partitioning: + return list(sql_commands) + return [command for command in sql_commands if not _is_spend_logs_pkey_rewrite(command)] + def extract_sql_commands(diff_output: str) -> list[str]: """ @@ -44,7 +75,9 @@ def extract_sql_commands(diff_output: str) -> list[str]: return sql_commands -def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: +def check_prisma_schema_diff_helper( + db_url: str, use_spend_logs_partitioning: bool = False +) -> tuple[bool, list[str]]: # mutable-ok: command list mirrors extract_sql_commands """Checks for differences between current database and Prisma schema. Never raises: a diff that cannot be produced, because the runner is missing, @@ -86,7 +119,9 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: env=os.environ.copy(), ) - sql_commands: Final = extract_sql_commands(result.stdout) + sql_commands: Final = filter_spend_logs_pkey_rewrite( + extract_sql_commands(result.stdout), use_spend_logs_partitioning + ) if sql_commands: print("Changes to DB Schema detected") # noqa: T201 @@ -108,13 +143,13 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: return False, [] -def check_prisma_schema_diff(db_url: str | None = None) -> None: +def check_prisma_schema_diff(db_url: str | None = None, use_spend_logs_partitioning: bool = False) -> None: """Main function to run the Prisma schema diff check.""" if db_url is None: db_url = os.getenv("DATABASE_URL") if db_url is None: 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, use_spend_logs_partitioning) if has_diff: verbose_logger.exception( "🚨🚨🚨 prisma schema out of sync with db. Consider running these sql_commands to sync the two - %s", diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 7f68b72a522..c569451297f 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -9,7 +9,7 @@ import sys import urllib.parse as urlparse from collections.abc import Iterable, Mapping, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, cast # noqa: TID251 # typed view of YAML settings import click import httpx @@ -1486,7 +1486,11 @@ def run_server( ) if should_update_prisma_schema(general_settings.get("disable_prisma_schema_update")) is False: - check_prisma_schema_diff(db_url=None) + settings_map: Final = cast(Mapping[str, object], general_settings) # cast-ok: YAML str-keyed map + check_prisma_schema_diff( + db_url=None, + use_spend_logs_partitioning=bool(settings_map.get("use_spend_logs_partitioning", False)), + ) else: use_v2_resolver: Final = resolve_v2_migration_resolver( use_legacy_flag=use_legacy_migration_resolver, diff --git a/tests/unit/proxy/db/test_check_migration.py b/tests/unit/proxy/db/test_check_migration.py index 74a6841cb23..bb5d17340b9 100644 --- a/tests/unit/proxy/db/test_check_migration.py +++ b/tests/unit/proxy/db/test_check_migration.py @@ -1,8 +1,6 @@ - import pytest - def test_check_migration_out_of_sync(mocker): """ Test that the check_prisma_schema_diff function @@ -51,8 +49,15 @@ def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fak assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) assert fake_prisma_cli.calls == [ - ["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x", - "--to-schema-datamodel", "./schema.prisma", "--script"] + [ + "migrate", + "diff", + "--from-url", + "postgresql://u:p@localhost:9/x", + "--to-schema-datamodel", + "./schema.prisma", + "--script", + ] ] assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) @@ -69,3 +74,75 @@ def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(m monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + + +PKEY_REWRITE = ( + 'ALTER TABLE "LiteLLM_SpendLogs" DROP CONSTRAINT "LiteLLM_SpendLogs_pkey", ' + 'ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id");' +) + + +def test_spend_logs_pkey_rewrite_filtered_when_partitioning_on(): + """Partitioned deployments (PK must include "startTime") must not be told + to rewrite the primary key to request_id alone - the rewrite is invalid + for the partitioned layout and schema.prisma's single-column @id always + produces it.""" + from litellm.proxy.db.check_migration import filter_spend_logs_pkey_rewrite + + assert filter_spend_logs_pkey_rewrite([PKEY_REWRITE], use_spend_logs_partitioning=True) == [] + + +def test_spend_logs_pkey_rewrite_kept_when_partitioning_off(): + """Without partitioning the rewrite is valid advice and must stay.""" + from litellm.proxy.db.check_migration import filter_spend_logs_pkey_rewrite + + assert filter_spend_logs_pkey_rewrite([PKEY_REWRITE], use_spend_logs_partitioning=False) == [PKEY_REWRITE] + + +def test_spend_logs_pkey_rewrite_filtered_in_two_statement_form(): + """The diff can also arrive as separate DROP / ADD statements.""" + from litellm.proxy.db.check_migration import filter_spend_logs_pkey_rewrite + + commands = [ + 'ALTER TABLE "LiteLLM_SpendLogs" DROP CONSTRAINT "LiteLLM_SpendLogs_pkey";', + 'ALTER TABLE "LiteLLM_SpendLogs" ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id");', + ] + + assert filter_spend_logs_pkey_rewrite(commands, use_spend_logs_partitioning=True) == [] + + +def test_real_drift_and_other_spend_logs_tables_survive_the_filter(): + """Unrelated drift, the legacy table, and partition child tables are not + touched by the filter - only the exact table's pkey rewrite is.""" + from litellm.proxy.db.check_migration import filter_spend_logs_pkey_rewrite + + drift = 'ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "new_field" TEXT;' + spend_logs_drift = 'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "new_field" TEXT;' + legacy = 'DROP TABLE "LiteLLM_SpendLogs_legacy";' + child = 'ALTER TABLE "LiteLLM_SpendLogs_2026_10" ADD CONSTRAINT "LiteLLM_SpendLogs_2026_10_pkey" PRIMARY KEY ("request_id", "startTime");' + commands = [PKEY_REWRITE, drift, spend_logs_drift, legacy, child] + + assert filter_spend_logs_pkey_rewrite(commands, use_spend_logs_partitioning=True) == [ + drift, + spend_logs_drift, + legacy, + child, + ] + + +def test_check_prisma_schema_diff_forwards_partitioning_flag( # test-quality-ok: forwarding test; the helper behaviour is covered by the filtering tests above + mocker, +): + """The boot-time wrapper must forward the flag to the helper.""" + from litellm.proxy.db import check_migration + + mocker.patch.object(check_migration, "verbose_logger", autospec=True) + helper = mocker.patch.object( + check_migration, + "check_prisma_schema_diff_helper", + return_value=(False, []), + ) + + check_migration.check_prisma_schema_diff(db_url="mock_url", use_spend_logs_partitioning=True) + + helper.assert_called_once_with("mock_url", True)