From 066c4f17cc0d665075f73068e40d1116de409700 Mon Sep 17 00:00:00 2001 From: Delweng Date: Mon, 12 Oct 2026 14:27:32 +0800 Subject: [PATCH 1/3] test(proxy): pin partition-aware startup schema advice for SpendLogs A partitioned LiteLLM_SpendLogs keeps "startTime" in its primary key, so the startup schema check must not report the ("request_id") primary-key rewrite as required SQL, while genuine drift on other tables in the same diff script still surfaces. The partition probe failing must degrade to today's advice instead of breaking boot. Signed-off-by: Delweng --- tests/unit/proxy/db/test_check_migration.py | 97 +++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/tests/unit/proxy/db/test_check_migration.py b/tests/unit/proxy/db/test_check_migration.py index 74a6841cb23..95aa8bd72a2 100644 --- a/tests/unit/proxy/db/test_check_migration.py +++ b/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,95 @@ 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\n" + 'ALTER 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] From 3e480a74624e82813b787d3737e79c84f2327299 Mon Sep 17 00:00:00 2001 From: Delweng Date: Mon, 12 Oct 2026 15:35:25 +0800 Subject: [PATCH 2/3] fix(proxy): skip partitioned SpendLogs primary-key drift in the startup schema check When LiteLLM_SpendLogs is partitioned by "startTime", prisma migrate diff reports a primary-key rewrite to ("request_id") that Postgres rejects (a partitioned table's key must include the partition key), and the startup schema check printed it as required SQL. The startup path now filters the already-extracted statements through litellm_proxy_extras' SpendLogs primary-key clause stripper, so unrelated drift in the same script stays reportable; chaining the raw-script filter before extraction would hide every statement, because extraction depends on the -- comment headers the filter strips. The clause split is now paren-aware and separator-tolerant, so it works on both Prisma's raw ',\n' clauses and the single-line ', ' form extraction produces. The partition probe is lazily imported and only runs once a diff was found, and probe failures (extras missing, psycopg missing, psycopg errors) degrade to the unfiltered advice with a skip note instead of breaking boot. Signed-off-by: Delweng --- .../litellm_proxy_extras/utils.py | 26 ++++-- litellm/proxy/db/check_migration.py | 84 +++++++++++++++++-- .../test_litellm_proxy_extras_utils.py | 18 ++++ 3 files changed, 117 insertions(+), 11 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 061fb627c65..26f4c2a28bb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -164,14 +164,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 @@ -189,7 +205,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) diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index 6e2a06c96e1..718f98ee1e7 100644 --- a/litellm/proxy/db/check_migration.py +++ b/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 diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 697c5a828bb..4e84376c4fe 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -24,6 +24,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 @@ -538,6 +539,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 = "" From 87e947cc8c6ee33fa19e75bf5e7ae22019dfccfe Mon Sep 17 00:00:00 2001 From: Delweng Date: Mon, 12 Oct 2026 15:35:49 +0800 Subject: [PATCH 3/3] style(test): collapse test constants and helper signature to one line Pure formatting on the new test block, matching ruff format's preference for lines that fit within 120 characters. Signed-off-by: Delweng --- tests/unit/proxy/db/test_check_migration.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/unit/proxy/db/test_check_migration.py b/tests/unit/proxy/db/test_check_migration.py index 95aa8bd72a2..9277774e8f3 100644 --- a/tests/unit/proxy/db/test_check_migration.py +++ b/tests/unit/proxy/db/test_check_migration.py @@ -82,15 +82,10 @@ _PARTITIONED_SPEND_LOGS_PK_DIFF: Final = ( 'ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id");\n' ) -_BUDGET_TABLE_DIFF: Final = ( - "-- AlterTable\n" - 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_at" TIMESTAMP(3);\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]]: +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),