diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index be785e2fd78..c2483617819 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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) 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 02bdf3429ea..1948eb4437b 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 @@ -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 = "" diff --git a/tests/unit/proxy/db/test_check_migration.py b/tests/unit/proxy/db/test_check_migration.py index 74a6841cb23..9277774e8f3 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,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]