scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix: guard accounting event removal across schema update paths

Browse source
This commit is contained in:
Joshua Valluru 2026-11-03 01:40:44 -08:00
parent e9f26e79cc
commit 10d37161c5
8 changed files with 208 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_remove_accounting_events/migration.sql
View file

@ -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;

32
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -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(),

1
litellm/integrations/otel/model/spans.py
View file

@ -339,7 +339,6 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_TelemetryReport",
"LiteLLM_AccessGroupTable",
"LiteLLM_AccountingDelivery",
"LiteLLM_AccountingEvent",
"LiteLLM_ClaudeCodePluginTable",
"LiteLLM_ClaudeCodePluginUpload",
"LiteLLM_MemoryTable",

1
litellm/proxy/db/prisma_client.py
View file

@ -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(
[

11
tests/code_coverage_tests/check_migrations_lock_timeout.py
View file

@ -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}",
)

104
tests/integration/database/test_spendlogs_projection.py
View file

@ -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()

49
tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py
View file

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

12
tests/unit/test_check_migrations_lock_timeout.py
View file

@ -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;")