From 10d37161c5ba8eaaa8caad6f8fcd48640578513a Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:40:44 -0800 Subject: [PATCH] fix: guard accounting event removal across schema update paths --- .../migration.sql | 12 +- .../litellm_proxy_extras/utils.py | 32 ++++++ litellm/integrations/otel/model/spans.py | 1 - litellm/proxy/db/prisma_client.py | 1 + .../check_migrations_lock_timeout.py | 11 +- .../database/test_spendlogs_projection.py | 104 ++++++++++++++++++ .../test_litellm_proxy_extras_utils.py | 49 ++++++++- .../test_check_migrations_lock_timeout.py | 12 ++ 8 files changed, 208 insertions(+), 14 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_remove_accounting_events/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_remove_accounting_events/migration.sql index a0fa06a300c..1e1712fa667 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_remove_accounting_events/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_remove_accounting_events/migration.sql @@ -1,12 +1,14 @@ BEGIN; SET LOCAL lock_timeout = '10s'; -LOCK TABLE "LiteLLM_AccountingEvent", "LiteLLM_AccountingDelivery" IN ACCESS EXCLUSIVE MODE; DO $$ BEGIN - IF EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" WHERE state <> 'applied') THEN - RAISE EXCEPTION 'Drain pending and parked accounting events before removing the event ledger'; + IF to_regclass('"LiteLLM_AccountingEvent"') IS NOT NULL THEN + LOCK TABLE "LiteLLM_AccountingEvent", "LiteLLM_AccountingDelivery" IN ACCESS EXCLUSIVE MODE; + IF EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" WHERE state <> 'applied') THEN + RAISE EXCEPTION 'Drain pending and parked accounting events before removing the event ledger'; + END IF; + DROP TABLE "LiteLLM_AccountingEvent"; END IF; END $$; -DROP TABLE "LiteLLM_AccountingEvent"; -DROP TYPE "AccountingEventKind"; +DROP TYPE IF EXISTS "AccountingEventKind"; COMMIT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 8b19576d86e..89cbfe53656 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -674,6 +674,36 @@ class ProxyExtrasDBManager: "Deployments using migration history can upgrade without --use_prisma_db_push instead." ) + @staticmethod + def raise_if_accounting_event_removal_pending() -> None: + database_url: Final = os.environ.get("DATABASE_URL") + if not database_url: + return + try: + import psycopg + except ImportError as exc: + raise RuntimeError("Install psycopg to verify accounting events before prisma db push.") from exc + try: + with psycopg.connect( + ProxyExtrasDBManager._strip_prisma_query_params(database_url), connect_timeout=10, autocommit=True + ) as connection: + events: Final = connection.execute( + "SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid=c.relnamespace " + "WHERE n.nspname=%s AND c.relname='LiteLLM_AccountingEvent' LIMIT 1", + (ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",), + ).fetchone() + except psycopg.Error as exc: + raise RuntimeError( + "Cannot verify accounting events; refusing prisma db push. " + f"The database check failed: {_redact_credentials(str(exc)).strip()}" + ) from exc + if events is not None: + raise RuntimeError( + "Legacy accounting events require the guarded removal migration. " + "Drain pending and parked deliveries, then upgrade without --use_prisma_db_push. " + "prisma db push bypasses the migration and can discard unfinished work." + ) + @staticmethod def spend_logs_is_partitioned() -> bool: """True when the connected database's LiteLLM_SpendLogs is a @@ -1037,6 +1067,7 @@ class ProxyExtrasDBManager: if not use_migrate: ProxyExtrasDBManager.raise_if_lens_rename_pending() + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() if ProxyExtrasDBManager.spend_logs_is_partitioned(): raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR) original_dir = os.getcwd() @@ -1715,6 +1746,7 @@ class ProxyExtrasDBManager: raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR) # Use prisma db push with increased timeout ProxyExtrasDBManager.raise_if_lens_rename_pending() + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() prisma_toolchain.run_prisma( [_get_prisma_command(), "db", "push", "--accept-data-loss"], timeout=prisma_command_timeout(), diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index ece85b09432..dda5de632b7 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -339,7 +339,6 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_TelemetryReport", "LiteLLM_AccessGroupTable", "LiteLLM_AccountingDelivery", - "LiteLLM_AccountingEvent", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginUpload", "LiteLLM_MemoryTable", diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index be5de040fdf..72d8d81ba75 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -1085,6 +1085,7 @@ class PrismaManager: from litellm_proxy_extras.utils import ProxyExtrasDBManager ProxyExtrasDBManager.raise_if_lens_rename_pending() + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() PrismaManager._raise_if_partitioned_spend_logs() run_prisma( [ diff --git a/tests/code_coverage_tests/check_migrations_lock_timeout.py b/tests/code_coverage_tests/check_migrations_lock_timeout.py index 22190f5fb77..20f8c978bb3 100644 --- a/tests/code_coverage_tests/check_migrations_lock_timeout.py +++ b/tests/code_coverage_tests/check_migrations_lock_timeout.py @@ -16,12 +16,13 @@ LOCK_TIMEOUT_RE: Final = re.compile( ) -def first_statement(sql: str) -> str: +def first_transaction_statement(sql: str) -> str: executable_sql: Final = "\n".join( line for line in sql.splitlines() if line.strip() and not line.lstrip().startswith("--") ) - match: Final = re.search(r"(?s).*?;", executable_sql) - return match.group().strip() if match is not None else executable_sql.strip() + transaction_sql: Final = re.sub(r"\ABEGIN\s*;\s*", "", executable_sql.strip(), count=1, flags=re.IGNORECASE) + match: Final = re.search(r"(?s).*?;", transaction_sql) + return match.group().strip() if match is not None else transaction_sql.strip() def scan_migration(directory: Path) -> tuple[str, ...]: @@ -32,11 +33,11 @@ def scan_migration(directory: Path) -> tuple[str, ...]: if "CONCURRENTLY" in sql.upper(): return () - if LOCK_TIMEOUT_RE.fullmatch(first_statement(sql)): + if LOCK_TIMEOUT_RE.fullmatch(first_transaction_statement(sql)): return () return ( - f"{directory.name}/migration.sql: first SQL statement must be SET LOCAL lock_timeout; " + f"{directory.name}/migration.sql: first SQL statement after an optional BEGIN must be SET LOCAL lock_timeout; " f"add this exact line: {LOCK_TIMEOUT_LINE}", ) diff --git a/tests/integration/database/test_spendlogs_projection.py b/tests/integration/database/test_spendlogs_projection.py index d3b42129c1f..f128a6d9b3b 100644 --- a/tests/integration/database/test_spendlogs_projection.py +++ b/tests/integration/database/test_spendlogs_projection.py @@ -49,3 +49,107 @@ def test_event_removal_preserves_applied_delivery_receipts( assert event_database.execute('SELECT event_id, state::text FROM "LiteLLM_AccountingDelivery"').fetchall() == [ ("event", "applied") ] + + +@pytest.mark.parametrize("entrypoint", ["proxy", "extras_v1", "extras_v2"]) +def test_schema_push_cannot_bypass_unfinished_event_guard( + event_database: psycopg.Connection[tuple[object, ...]], monkeypatch: pytest.MonkeyPatch, entrypoint: str +) -> None: + import os + from urllib.parse import urlsplit, urlunsplit + from litellm.proxy.db.prisma_client import PrismaManager + from litellm_proxy_extras.utils import ProxyExtrasDBManager + + event_database.execute( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state) ' + "VALUES ('event', 'LiteLLM_UserTable', 1, 'pending')" + ) + url = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{event_database.info.dbname}")) + monkeypatch.setenv("DATABASE_URL", url) + setup_database = PrismaManager.setup_database if entrypoint == "proxy" else ProxyExtrasDBManager.setup_database + with pytest.raises(RuntimeError, match="accounting events"): + setup_database(use_migrate=False, use_v2_resolver=entrypoint == "extras_v2") + assert event_database.execute('SELECT event_id FROM "LiteLLM_AccountingEvent"').fetchall() == [("event",)] + assert event_database.execute('SELECT state::text FROM "LiteLLM_AccountingDelivery"').fetchall() == [("pending",)] + + +@pytest.mark.parametrize("interrupt_before_commit", [False, True]) +def test_prisma_deploy_preserves_event_removal_transaction( + event_database: psycopg.Connection[tuple[object, ...]], tmp_path: Path, interrupt_before_commit: bool +) -> None: + import os + import subprocess + import sys + from urllib.parse import urlsplit, urlunsplit + + url = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{event_database.info.dbname}")) + schema = tmp_path / "schema.prisma" + schema.write_text('datasource db {\n provider = "postgresql"\n url = env("DATABASE_URL")\n}\n') + migrations = tmp_path / "migrations" + migrations.mkdir() + (migrations / "migration_lock.toml").write_text('provider = "postgresql"\n') + initial_name = "20261027000000_add_accounting_ledger_tables" + candidate_name = "20261103000000_remove_accounting_events" + for name in (initial_name, candidate_name): + directory = migrations / name + directory.mkdir() + statement = (_MIGRATIONS / name / "migration.sql").read_text() + if name == candidate_name and interrupt_before_commit: + statement = statement.replace("COMMIT;", "SELECT 1 / 0;\nCOMMIT;") + (directory / "migration.sql").write_text(statement) + command = [sys.executable, "-I", "-m", "prisma", "migrate"] + environment = {**os.environ, "DATABASE_URL": url} + resolved = subprocess.run( + command + ["resolve", "--applied", initial_name, "--schema", str(schema)], + env=environment, + capture_output=True, + text=True, + ) + assert resolved.returncode == 0, resolved.stderr + result = subprocess.run( + command + ["deploy", "--schema", str(schema)], + env=environment, + capture_output=True, + text=True, + ) + present = event_database.execute("SELECT to_regclass('\"LiteLLM_AccountingEvent\"') IS NOT NULL").fetchone() + if interrupt_before_commit: + assert result.returncode != 0, result.stdout + assert present == (True,) + assert event_database.execute('SELECT event_id FROM "LiteLLM_AccountingEvent"').fetchall() == [("event",)] + else: + assert result.returncode == 0, result.stderr + assert present == (False,) + assert event_database.execute( + 'SELECT finished_at IS NOT NULL FROM "_prisma_migrations" WHERE migration_name=%s', (candidate_name,) + ).fetchone() == (True,) + + +def test_committed_removal_can_replay_after_lost_migration_record( + event_database: psycopg.Connection[tuple[object, ...]], +) -> None: + statement = (_MIGRATIONS / "20261103000000_remove_accounting_events/migration.sql").read_text() + event_database.execute(statement) + event_database.execute(statement) + assert event_database.execute("SELECT to_regclass('\"LiteLLM_AccountingEvent\"')").fetchone() == (None,) + assert event_database.execute('SELECT COUNT(*) FROM "LiteLLM_AccountingDelivery"').fetchone() == (0,) + + +def test_schema_push_guard_only_blocks_events_in_the_target_schema( + event_database: psycopg.Connection[tuple[object, ...]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + import os + from urllib.parse import urlsplit, urlunsplit + from litellm_proxy_extras.utils import ProxyExtrasDBManager + + url = urlunsplit( + urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{event_database.info.dbname}", query="schema=other") + ) + monkeypatch.setenv("DATABASE_URL", url) + event_database.execute("CREATE SCHEMA other") + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() + assert event_database.execute('SELECT event_id FROM "LiteLLM_AccountingEvent"').fetchall() == [("event",)] + event_database.execute('CREATE TABLE other."LiteLLM_AccountingEvent" (event_id text)') + with pytest.raises(RuntimeError, match="accounting events"): + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() 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..96fa8d758c6 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 @@ -753,13 +753,17 @@ class TestPartitionedSpendLogsPushGuard: class _FakeCursor: + def __init__(self, present=True): + self._present = present + def fetchone(self): - return (1,) + return (1,) if self._present else None class _FakePsycopgConn: - def __init__(self, executed): + def __init__(self, executed, present=True): self._executed = executed + self._present = present def __enter__(self): return self @@ -769,7 +773,7 @@ class _FakePsycopgConn: def execute(self, query, params): self._executed.append((query, params)) - return _FakeCursor() + return _FakeCursor(self._present) class TestLensRenamePendingCheck: @@ -2373,3 +2377,42 @@ class TestMigrateDeployRunsUnderTheCoordinatorLock: assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True assert events == ["lock", "deploy", "unlock", "lock", "rollback", "unlock", "lock", "deploy", "unlock"] + + +class TestAccountingEventRemovalPushGuard: + def test_missing_database_url_needs_no_schema_check(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DATABASE_URL", raising=False) + assert ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() is None + + def test_missing_database_driver_refuses_schema_push(self, monkeypatch: pytest.MonkeyPatch) -> None: + import sys + + monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/database") + monkeypatch.setitem(sys.modules, "psycopg", None) + with pytest.raises(RuntimeError, match="Install psycopg"): + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() + + def test_failed_schema_check_redacts_database_credentials(self, monkeypatch: pytest.MonkeyPatch) -> None: + password: Final = "p%zzword" + monkeypatch.setenv("DATABASE_URL", f"postgresql://litellm:{password}@localhost:9/database") + monkeypatch.delenv("DIRECT_URL", raising=False) + with pytest.raises(RuntimeError, match="Cannot verify accounting events") as error: + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() + assert password not in str(error.value) + + + @pytest.mark.parametrize("present", [True, False]) + def test_existing_events_require_migration_instead_of_schema_push( + self, monkeypatch: pytest.MonkeyPatch, present: bool, + ) -> None: + import psycopg + + monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/database?schema=accounting") + executed = [] + monkeypatch.setattr(psycopg, "connect", lambda *args, **kwargs: _FakePsycopgConn(executed, present)) + if present: + with pytest.raises(RuntimeError, match="guarded removal migration"): + ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() + else: + assert ProxyExtrasDBManager.raise_if_accounting_event_removal_pending() is None + assert len(executed) == 1 and executed[0][1] == ("accounting",) diff --git a/tests/unit/test_check_migrations_lock_timeout.py b/tests/unit/test_check_migrations_lock_timeout.py index b9c2f723fba..d94f36b8c1f 100644 --- a/tests/unit/test_check_migrations_lock_timeout.py +++ b/tests/unit/test_check_migrations_lock_timeout.py @@ -60,3 +60,15 @@ def test_concurrently_migration_without_lock_timeout_passes(tmp_path: Path) -> N sql = "create index concurrently foo_bar_idx on foo (bar);" assert _scan(tmp_path, sql) == () + + +def test_transaction_sets_lock_timeout_before_acquiring_locks(tmp_path: Path) -> None: + assert _scan(tmp_path, "BEGIN;\nSET LOCAL lock_timeout = '10s';\nLOCK TABLE foo;\nCOMMIT;") == () + + +def test_transaction_without_lock_timeout_is_rejected(tmp_path: Path) -> None: + assert _scan(tmp_path, "BEGIN;\nLOCK TABLE foo;\nCOMMIT;") + + +def test_transaction_cannot_set_lock_timeout_after_acquiring_locks(tmp_path: Path) -> None: + assert _scan(tmp_path, "BEGIN;\nLOCK TABLE foo;\nSET LOCAL lock_timeout = '10s';\nCOMMIT;")