Merge 87e947cc8c into d56ff281de

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -163,14 +163,30 @@ def _without_sql_comments(statement: str) -> str:
     ).strip()
 
 
-def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
+_TOP_LEVEL_COMMA_RE = re.compile(r"[^,()]+(?:\([^()]*\))?[^,()]*")
+
+
+def _split_top_level_commas(text: str) -> tuple[str, ...]:
+    """Comma-separated pieces of `text`, ignoring commas inside parentheses
+    (the column list of `PRIMARY KEY ("request_id", "startTime")`). A comma
+    that lands inside a SQL literal stays inside its piece, so the pieces
+    rejoin to the same SQL."""
+    return tuple(piece.strip() for piece in _TOP_LEVEL_COMMA_RE.findall(text) if piece.strip())
+
+
+def without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
+    """One ALTER TABLE statement (no trailing semicolon) without the
+    LiteLLM_SpendLogs primary-key-rewrite clauses, or None when nothing
+    remains. Clause separators may be Prisma's raw `,\n` or the single-line
+    `, ` form `extract_sql_commands` produces, so both the raw drift script
+    and already-extracted statements can be filtered."""
     prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
     if not prefix_match:
         return statement
     kept = tuple(
-        clause.strip()
-        for clause in statement[prefix_match.end():].split(",\n")
-        if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause.strip())
+        clause
+        for clause in _split_top_level_commas(statement[prefix_match.end():])
+        if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause)
     )
     if not kept:
         return None
@@ -188,7 +204,7 @@ def filter_partitioned_spend_logs_diff(diff_sql: str) -> str:
         for statement in diff_sql.split(";")
         for bare in (_without_sql_comments(statement),)
         if bare and not _SPEND_LOGS_ARTIFACT_DROP_RE.match(bare)
-        for filtered in (_without_spend_logs_pk_clauses(bare),)
+        for filtered in (without_spend_logs_pk_clauses(bare),)
         if filtered is not None
     )
     return "".join(f"{statement};\n\n" for statement in kept)

litellm/proxy/db/check_migration.py

@@ -2,6 +2,7 @@
 
 import os
 import subprocess
+from collections.abc import Callable, Iterator, Sequence
 from typing import Final
 
 from litellm._logging import verbose_logger
@@ -44,17 +45,85 @@ 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 _spend_logs_is_partitioned() -> bool:
+    """The litellm-proxy-extras partition probe, degraded to False when it cannot run."""
+    try:
+        from litellm_proxy_extras.utils import (
+            ProxyExtrasDBManager,  # pyright: ignore[reportUnknownVariableType]  # optional untyped install
+        )
+    except ImportError as e:
+        verbose_logger.warning(
+            "Skipping the partitioned LiteLLM_SpendLogs check: litellm-proxy-extras is unavailable. Error: %s", e
+        )
+        return False
+    try:
+        import psycopg
+    except ImportError:
+        verbose_logger.warning(
+            "psycopg is not installed, so the partitioned LiteLLM_SpendLogs check was skipped and the SQL "
+            "below may include a primary-key rewrite Postgres rejects on a partitioned table"
+        )
+        return False
+    try:
+        return ProxyExtrasDBManager.spend_logs_is_partitioned()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # optional untyped install
+    except psycopg.Error as e:
+        verbose_logger.warning(
+            "The partitioned LiteLLM_SpendLogs check failed, so the SQL below may include a primary-key "
+            "rewrite Postgres rejects on a partitioned table. Error: %s",
+            e,
+        )
+        return False
+
+
+def _statements_without_spend_logs_pk_rewrite(
+    sql_commands: Sequence[str], strip_pk_clauses: Callable[[str], str | None]
+) -> Iterator[str]:
+    for statement in sql_commands:
+        without_pk = strip_pk_clauses(statement.strip().removesuffix(";"))
+        if without_pk is not None:
+            yield f"{without_pk};" if statement.endswith(";") else without_pk
+
+
+def _without_partitioned_spend_logs_pk_rewrite(
+    sql_commands: Sequence[str], *, partitioned: bool | None
+) -> tuple[str, ...]:
+    """The extracted statements of a diff script, minus the LiteLLM_SpendLogs
+    primary-key rewrite when the table is partitioned: Postgres rejects it there
+    (the key must include the partition key), so it is not executable advice.
+    Filtering the already-extracted statements keeps unrelated drift reportable."""
+    if not sql_commands:
+        return tuple(sql_commands)
+    is_partitioned: Final = partitioned if partitioned is not None else _spend_logs_is_partitioned()
+    if not is_partitioned:
+        return tuple(sql_commands)
+    try:
+        from litellm_proxy_extras.utils import (
+            without_spend_logs_pk_clauses,  # pyright: ignore[reportUnknownVariableType]  # optional untyped install
+        )
+    except ImportError as e:
+        verbose_logger.warning(
+            "Skipping the partitioned LiteLLM_SpendLogs primary-key filter: "
+            "litellm-proxy-extras is unavailable. Error: %s",
+            e,
+        )
+        return tuple(sql_commands)
+    return tuple(_statements_without_spend_logs_pk_rewrite(sql_commands, without_spend_logs_pk_clauses))  # pyright: ignore[reportUnknownArgumentType]  # optional untyped install
+
+
+def check_prisma_schema_diff_helper(db_url: str, *, partitioned: bool | None = None) -> tuple[bool, list[str]]:
     """Checks for differences between current database and Prisma schema.
 
     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.
+    Args:
+        db_url: The database URL to diff against the Prisma schema.
+        partitioned: Whether LiteLLM_SpendLogs is partitioned. None resolves it
+            from the database, but only once a diff was found, since the probe
+            opens a second connection. On True, the primary-key rewrite Postgres
+            rejects on a partitioned LiteLLM_SpendLogs is left out of the
+            reported SQL.
     """
     try:
         from litellm_proxy_extras.prisma_toolchain import (
@@ -86,7 +155,10 @@ 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)
+        extracted_commands: Final = extract_sql_commands(result.stdout)
+        sql_commands: Final = list(
+            _without_partitioned_spend_logs_pk_rewrite(extracted_commands, partitioned=partitioned)
+        )
 
         if sql_commands:
             print("Changes to DB Schema detected")  # noqa: T201

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -25,6 +25,7 @@ from litellm_proxy_extras.utils import (
     _redact_command_error,
     _redact_credentials,
     filter_partitioned_spend_logs_diff,
+    without_spend_logs_pk_clauses,
 )
 
 # Path to the migrations directory
@@ -655,6 +656,23 @@ class TestPartitionedSpendLogsDriftFilter:
         assert 'DROP CONSTRAINT "LiteLLM_TeamTable_pkey"' in filtered
         assert 'PRIMARY KEY ("team_id")' in filtered
 
+    def test_a_single_line_statement_loses_only_the_pk_clause_and_keeps_the_rest(self):
+        statement = (
+            'ALTER TABLE "LiteLLM_SpendLogs" DROP CONSTRAINT "LiteLLM_SpendLogs_pkey", '
+            'ADD COLUMN     "end_time" TIMESTAMP(3), '
+            'ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id", "startTime")'
+        )
+        assert without_spend_logs_pk_clauses(statement) == (
+            'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN     "end_time" TIMESTAMP(3)'
+        )
+
+    def test_a_single_line_statement_with_only_the_pk_rewrite_is_dropped_entirely(self):
+        statement = (
+            'ALTER TABLE "LiteLLM_SpendLogs" DROP CONSTRAINT "LiteLLM_SpendLogs_pkey", '
+            'ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id")'
+        )
+        assert without_spend_logs_pk_clauses(statement) is None
+
 
 class _FakeCompleted:
     stdout = ""

tests/unit/proxy/db/test_check_migration.py

@@ -1,5 +1,10 @@
 
+import subprocess
+import sys
+from typing import Final
+
 import pytest
+from pytest_mock import MockerFixture
 
 
 
@@ -69,3 +74,90 @@ 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, [])
+
+
+_PARTITIONED_SPEND_LOGS_PK_DIFF: Final = (
+    "-- AlterTable\n"
+    'ALTER TABLE "LiteLLM_SpendLogs" DROP CONSTRAINT "LiteLLM_SpendLogs_pkey",\n'
+    'ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id");\n'
+)
+
+_BUDGET_TABLE_DIFF: Final = '-- AlterTable\nALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_at" TIMESTAMP(3);\n'
+
+
+def _schema_diff_result(mocker: MockerFixture, diff_sql: str, *, partitioned: bool | None) -> tuple[bool, list[str]]:
+    mocker.patch(
+        "litellm_proxy_extras.prisma_toolchain.run_prisma",
+        return_value=subprocess.CompletedProcess(args=[], returncode=0, stdout=diff_sql),
+    )
+    from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
+
+    return check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x", partitioned=partitioned)
+
+
+def test_partitioned_spend_logs_pk_rewrite_alone_reports_the_schema_in_sync(mocker: MockerFixture) -> None:
+    """
+    A partitioned LiteLLM_SpendLogs keeps "startTime" in its primary key, so the primary-key
+    rewrite Prisma's diff suggests is not executable SQL and is not advice at all.
+    """
+    assert _schema_diff_result(mocker, _PARTITIONED_SPEND_LOGS_PK_DIFF, partitioned=True) == (False, [])
+
+
+def test_partitioned_spend_logs_pk_filter_keeps_unrelated_tables_in_the_advice(mocker: MockerFixture) -> None:
+    has_diff, sql_commands = _schema_diff_result(
+        mocker, _PARTITIONED_SPEND_LOGS_PK_DIFF + _BUDGET_TABLE_DIFF, partitioned=True
+    )
+    assert has_diff is True
+    assert sql_commands == ['ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_at" TIMESTAMP(3);']
+
+
+def test_unpartitioned_spend_logs_still_reports_its_pk_drift(mocker: MockerFixture) -> None:
+    has_diff, sql_commands = _schema_diff_result(mocker, _PARTITIONED_SPEND_LOGS_PK_DIFF, partitioned=False)
+    assert has_diff is True
+    assert len(sql_commands) == 1
+    assert 'DROP CONSTRAINT "LiteLLM_SpendLogs_pkey"' in sql_commands[0]
+    assert 'PRIMARY KEY ("request_id")' in sql_commands[0]
+
+
+def test_a_failing_partition_probe_does_not_break_startup(mocker: MockerFixture) -> None:
+    """
+    The probe runs before uvicorn starts and psycopg's InterfaceError escapes its own
+    handler, so boot would die on it. It must degrade to the unfiltered advice instead.
+    """
+    import psycopg
+
+    from litellm_proxy_extras.utils import ProxyExtrasDBManager
+
+    mocker.patch.object(
+        ProxyExtrasDBManager,
+        "spend_logs_is_partitioned",
+        side_effect=psycopg.InterfaceError("another operation closed the connection"),
+    )
+    has_diff, sql_commands = _schema_diff_result(mocker, _PARTITIONED_SPEND_LOGS_PK_DIFF, partitioned=None)
+    assert has_diff is True
+    assert 'PRIMARY KEY ("request_id")' in sql_commands[0]
+
+
+def test_missing_psycopg_skips_the_partition_check_with_a_note_instead_of_bad_advice_silently(
+    mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from litellm.proxy.db import check_migration
+
+    monkeypatch.setitem(sys.modules, "psycopg", None)
+    mock_logger = mocker.patch.object(check_migration, "verbose_logger", autospec=True)
+
+    has_diff, sql_commands = _schema_diff_result(mocker, _PARTITIONED_SPEND_LOGS_PK_DIFF, partitioned=None)
+
+    assert has_diff is True
+    assert 'PRIMARY KEY ("request_id")' in sql_commands[0]
+    mock_logger.warning.assert_called_once()
+    assert "partitioned" in mock_logger.warning.call_args[0][0]
+
+
+def test_missing_proxy_extras_skips_the_partition_check_instead_of_crashing_boot(
+    mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None)
+    has_diff, sql_commands = _schema_diff_result(mocker, _PARTITIONED_SPEND_LOGS_PK_DIFF, partitioned=None)
+    assert has_diff is True
+    assert 'PRIMARY KEY ("request_id")' in sql_commands[0]