Merge dd12632688 into 748c973f42

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",

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,

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)