mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-30 01:52:18 +00:00
56 lines
2.3 KiB
Python
56 lines
2.3 KiB
Python
"""Shared, destructive DB helpers for the e2e harness.
|
|
|
|
Kept at the top level next to e2e_config and lifecycle so every suite imports it
|
|
by name (`from e2e_db import ...`); no suite reaches into another's directory by
|
|
mutating sys.path.
|
|
|
|
reset_spend_logs truncates LiteLLM_SpendLogs and cannot be undone, so the
|
|
session-finish cleanup routes through run_spend_log_cleanup, which fires the
|
|
truncate only on an explicit operator opt-in. "An e2e test ran" is necessary but
|
|
never sufficient: a DATABASE_URL pointing at a shared or staging instance must
|
|
not be wiped by a routine local run that merely exercised a test.
|
|
"""
|
|
|
|
import os
|
|
from collections.abc import Callable
|
|
|
|
RESET_OPT_IN_ENV = "E2E_RESET_SPEND_LOGS"
|
|
|
|
|
|
def run_spend_log_cleanup(
|
|
*, opt_in: str | None, e2e_test_ran: bool, truncate: Callable[[], None]
|
|
) -> bool:
|
|
"""Invoke `truncate` iff the destructive spend-log reset is both opted into
|
|
and warranted, returning whether the truncate was attempted.
|
|
|
|
The truncate fires only when the opt-in value is exactly "1" AND an e2e test
|
|
body actually ran. Any other opt-in value (unset, "0", "true", "") leaves the
|
|
DB untouched, so the destructive path is never armed by the env var's mere
|
|
presence or by a test run on its own. Best-effort: a truncate failure is
|
|
swallowed so cleanup never fails the session, so the returned bool reports
|
|
that the reset was attempted, not that the DB call succeeded.
|
|
"""
|
|
if opt_in != "1" or not e2e_test_ran:
|
|
return False
|
|
try:
|
|
truncate()
|
|
except Exception as exc: # noqa: BLE001 - cleanup is best-effort
|
|
print(f"spend-log cleanup best-effort failed: {exc}")
|
|
return True
|
|
|
|
|
|
def reset_spend_logs() -> None:
|
|
"""Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes
|
|
spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses
|
|
DATABASE_URL (default: the local docker postgres on its mapped host port; the
|
|
in-container `@db` host isn't resolvable from the host, so default to
|
|
localhost).
|
|
"""
|
|
import psycopg
|
|
|
|
url = os.environ.get(
|
|
"DATABASE_URL",
|
|
"postgresql://llmproxy:dbpassword9090@localhost:5432/litellm",
|
|
)
|
|
with psycopg.connect(url) as conn:
|
|
_ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"')
|