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

refactor(accounting): retain schema and prepare durable payloads

Browse source
This commit is contained in:
Joshua Valluru 2026-11-03 08:15:45 -08:00
parent 10d37161c5
commit c59e08e866
13 changed files with 268 additions and 281 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,14 +0,0 @@
BEGIN;
SET LOCAL lock_timeout = '10s';
DO $$
BEGIN
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 TYPE IF EXISTS "AccountingEventKind";
COMMIT;

23
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -2071,6 +2071,12 @@ model LiteLLM_TelemetryReport {
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
@ -2102,6 +2108,23 @@ enum AccountingDeliveryTarget {
LiteLLM_AutoRouterDailySpend
}
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {

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

@ -674,36 +674,6 @@ 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
@ -1067,7 +1037,6 @@ 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()
@ -1746,7 +1715,6 @@ 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(),

42
litellm/accounting/projection.py
View file

@ -68,15 +68,28 @@ def _spend(table: str, key: tuple[tuple[str, str], ...], spend: float) -> Contri
def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:
spend: Final = _number(row.get("spend"))
entities: Final = _fields(metadata["accounting_entities"]) if "accounting_entities" in metadata else row
identities: Final = (
("LiteLLM_VerificationToken", "token", row.get("api_key")),
("LiteLLM_UserTable", "user_id", row.get("user")),
("LiteLLM_UserTable", "user_id", metadata.get("accounting_proxy_budget_name")),
("LiteLLM_TeamTable", "team_id", row.get("team_id")),
("LiteLLM_OrganizationTable", "organization_id", row.get("organization_id")),
("LiteLLM_EndUserTable", "user_id", row.get("end_user")),
("LiteLLM_AgentsTable", "agent_id", row.get("agent_id")),
("LiteLLM_ProjectTable", "project_id", metadata.get("user_api_key_project_id")),
("LiteLLM_VerificationToken", "token", entities.get("api_key")),
("LiteLLM_UserTable", "user_id", entities.get("user")),
(
"LiteLLM_UserTable",
"user_id",
entities.get("proxy_budget_name")
if "accounting_entities" in metadata
else metadata.get("accounting_proxy_budget_name"),
),
("LiteLLM_TeamTable", "team_id", entities.get("team_id")),
("LiteLLM_OrganizationTable", "organization_id", entities.get("organization_id")),
("LiteLLM_EndUserTable", "user_id", entities.get("end_user")),
("LiteLLM_AgentsTable", "agent_id", entities.get("agent_id")),
(
"LiteLLM_ProjectTable",
"project_id",
entities.get("project_id")
if "accounting_entities" in metadata
else metadata.get("user_api_key_project_id"),
),
)
groups: Final = _STRINGS.validate_python(metadata.get("accounting_model_access_groups", ()))
return (
@ -87,7 +100,7 @@ def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, obje
),
*(_spend("LiteLLM_TagTable", (("tag_name", tag),), spend) for tag in _tags(row) if tag),
*(_spend("LiteLLM_ModelAccessGroupBudgetTable", (("access_group_name", group),), spend) for group in groups),
*_membership_contributions(row, spend),
*_membership_contributions(entities, spend),
)
@ -151,7 +164,14 @@ def _daily_contributions(row: Mapping[str, object], metadata: Mapping[str, objec
("mcp_namespaced_tool_name", _text(row.get("mcp_namespaced_tool_name"))),
("endpoint", _text(endpoint)),
)
return tuple(_daily_rows(row, shared_key, counters))
# The legacy daily organization writer uses the charged org argument, not
# SpendLogs' normalized empty-string identity for a request with no org.
daily_row: Final = (
{**row, "organization_id": _fields(metadata["accounting_entities"]).get("organization_id")}
if "accounting_entities" in metadata
else row
)
return tuple(_daily_rows(daily_row, shared_key, counters))
def _daily_rows(
@ -178,7 +198,7 @@ def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, obje
Contribution(
"LiteLLM_BudgetWindowSpend",
tuple((column, _text(window[column])) for column in ("entity_type", "entity_id", "window_duration")),
(("spend", _number(row.get("spend"))),),
(("spend", _number(window.get("spend", row.get("spend")))),),
(("window_start", _text(window["window_start"])), ("started_at", _text(row.get("startTime")))),
)
for window in windows

8
litellm/accounting/protocols.py
View file

@ -1,15 +1,13 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Protocol
from collections.abc import Mapping
from typing import Protocol
from litellm.accounting.types import AggregateQuery, AggregateRow
if TYPE_CHECKING:
from litellm.proxy._types import SpendLogsPayload
class CostAggregator(Protocol):
async def apply(self, rows: tuple[SpendLogsPayload, ...]) -> None:
async def apply(self, rows: tuple[Mapping[str, object], ...]) -> None:
"""Apply stored request contributions atomically with their delivery receipts."""
...

117
litellm/accounting/snapshot.py Normal file
View file

@ -0,0 +1,117 @@
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens
from litellm.proxy.spend_tracking.savings import compute_savings_spend, marks_gateway_injection
if TYPE_CHECKING:
from litellm.router import Router
_FIELDS: Final = TypeAdapter(dict[str, object])
def _mapping(value: object) -> Mapping[str, object] | None:
return _FIELDS.validate_python(value) if isinstance(value, Mapping) else None
def accounting_metadata(
payload: SpendLogsPayload,
*,
entities: Mapping[str, str | None],
model_access_groups: Sequence[str],
windows: Sequence[Mapping[str, object]],
router: "Callable[[], Router | None]",
) -> str:
metadata: Final = _FIELDS.validate_json(payload["metadata"])
savings: Final = compute_savings_spend(
model=payload["model"],
custom_llm_provider=payload["custom_llm_provider"],
compression_saved_tokens=extract_compression_saved_tokens(metadata),
gateway_injected_cache=marks_gateway_injection(metadata, payload["model_id"]),
routing_decision=_mapping(metadata.get("routing_decision")),
model_id=payload["model_id"],
llm_router=router,
usage_object=_mapping(metadata.get("usage_object")),
cost_breakdown=_mapping(metadata.get("cost_breakdown")),
recorded_autorouter_savings=metadata.get("autorouter_savings"),
recorded_autorouter_savings_estimate=_mapping(metadata.get("autorouter_savings_estimate")),
billed_at=payload["endTime"],
)
return safe_dumps(
{
**metadata,
"accounting_entities": entities,
"accounting_savings": savings._asdict(),
"accounting_model_access_groups": model_access_groups,
"accounting_windows": windows,
},
non_finite_as_null=True,
)
ACCOUNTING_METADATA_FIELDS: Final = frozenset(
{
"accounting_windows",
"accounting_lean",
"accounting_model_access_groups",
"accounting_savings",
"accounting_entities",
"accounting_proxy_budget_name",
"user_api_key_project_id",
"usage_object",
"status",
"guardrail_information",
"standard_logging_guardrail_information",
"compression_savings",
"internal_call_origin",
"billed_retry_attempt",
}
)
_ACCOUNTING_FIELDS: Final = frozenset(
{
"request_id",
"call_type",
"api_key",
"spend",
"prompt_tokens",
"completion_tokens",
"startTime",
"model",
"model_group",
"custom_llm_provider",
"user",
"team_id",
"organization_id",
"end_user",
"agent_id",
"request_tags",
"status",
"request_duration_ms",
"mcp_namespaced_tool_name",
}
)
def build_accounting_payload(
payload: SpendLogsPayload,
*,
entities: Mapping[str, str | None],
model_access_groups: Sequence[str],
windows: Sequence[Mapping[str, object]],
router: "Callable[[], Router | None]",
) -> Mapping[str, object]:
metadata: Final = _FIELDS.validate_json(
accounting_metadata(
payload, entities=entities, model_access_groups=model_access_groups, windows=windows, router=router
)
)
return {
**{key: value for key, value in payload.items() if key in _ACCOUNTING_FIELDS},
"metadata": {key: value for key, value in metadata.items() if key in ACCOUNTING_METADATA_FIELDS},
}

23
litellm/proxy/schema.prisma
View file

@ -2071,6 +2071,12 @@ model LiteLLM_TelemetryReport {
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
@ -2102,6 +2108,23 @@ enum AccountingDeliveryTarget {
LiteLLM_AutoRouterDailySpend
}
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {

23
schema.prisma
View file

@ -2071,6 +2071,12 @@ model LiteLLM_TelemetryReport {
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
@ -2102,6 +2108,23 @@ enum AccountingDeliveryTarget {
LiteLLM_AutoRouterDailySpend
}
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {

11
tests/code_coverage_tests/check_migrations_lock_timeout.py
View file

@ -16,13 +16,12 @@ LOCK_TIMEOUT_RE: Final = re.compile(
)
def first_transaction_statement(sql: str) -> str:
def first_statement(sql: str) -> str:
executable_sql: Final = "\n".join(
line for line in sql.splitlines() if line.strip() and not line.lstrip().startswith("--")
)
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()
match: Final = re.search(r"(?s).*?;", executable_sql)
return match.group().strip() if match is not None else executable_sql.strip()
def scan_migration(directory: Path) -> tuple[str, ...]:
@ -33,11 +32,11 @@ def scan_migration(directory: Path) -> tuple[str, ...]:
if "CONCURRENTLY" in sql.upper():
return ()
if LOCK_TIMEOUT_RE.fullmatch(first_transaction_statement(sql)):
if LOCK_TIMEOUT_RE.fullmatch(first_statement(sql)):
return ()
return (
f"{directory.name}/migration.sql: first SQL statement after an optional BEGIN must be SET LOCAL lock_timeout; "
f"{directory.name}/migration.sql: first SQL statement must be SET LOCAL lock_timeout; "
f"add this exact line: {LOCK_TIMEOUT_LINE}",
)

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

@ -1,155 +0,0 @@
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import psycopg
import pytest
from tests.integration._support.database import scratch_database
_MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/migrations"
@pytest.fixture
def event_database() -> Iterator[psycopg.Connection[tuple[object, ...]]]:
with scratch_database("spendlogs_projection") as url, psycopg.connect(url, autocommit=True) as connection:
connection.execute((_MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql").read_text())
connection.execute(
'INSERT INTO "LiteLLM_AccountingEvent" (event_id, kind, schema_version, payload_hash, occurred_at, payload) '
"VALUES ('event', 'outbound_usage', 1, 'hash', '2026-11-02', '{}')"
)
yield connection
@pytest.mark.parametrize("state", ["pending", "parked"])
def test_event_removal_refuses_unfinished_work(
event_database: psycopg.Connection[tuple[object, ...]], state: str
) -> None:
event_database.execute(
'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state) '
"VALUES ('event', 'LiteLLM_UserTable', 1, %s)",
(state,),
)
with pytest.raises(psycopg.errors.RaiseException, match="Drain pending and parked"):
event_database.execute((_MIGRATIONS / "20261103000000_remove_accounting_events/migration.sql").read_text())
event_database.execute("ROLLBACK")
assert event_database.execute('SELECT event_id FROM "LiteLLM_AccountingEvent"').fetchall() == [("event",)]
assert event_database.execute('SELECT state::text FROM "LiteLLM_AccountingDelivery"').fetchall() == [(state,)]
def test_event_removal_preserves_applied_delivery_receipts(
event_database: psycopg.Connection[tuple[object, ...]],
) -> None:
event_database.execute(
'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state) '
"VALUES ('event', 'LiteLLM_UserTable', 1, 'applied')"
)
event_database.execute((_MIGRATIONS / "20261103000000_remove_accounting_events/migration.sql").read_text())
assert event_database.execute("SELECT to_regclass('\"LiteLLM_AccountingEvent\"')").fetchone() == (None,)
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()

40
tests/unit/accounting/test_projection.py
View file

@ -7,7 +7,9 @@ import pytest
from pydantic import TypeAdapter
from litellm.accounting.projection import project
from litellm.accounting.snapshot import build_accounting_payload
from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.proxy_server import proxy_logging_obj
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
@ -51,6 +53,23 @@ async def test_daily_projection_matches_legacy_writer(status: str, origin: str |
assert dict(projected.increments) == {name: legacy[name] for name, _ in projected.increments}
assert dict(projected.key)["date"] == legacy["date"]
assert dict(projected.key)["api_key"] == legacy["api_key"]
accounting_input: Final[SpendLogsPayload] = {
**payload,
"spend": spend,
"metadata": json.dumps(metadata),
"request_duration_ms": 25,
}
accounting: Final = build_accounting_payload(
accounting_input,
entities={"api_key": payload["api_key"], "user": payload["user"]},
model_access_groups=(),
windows=(),
router=lambda: None,
)
restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str))
stored_daily: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend")
assert dict(stored_daily.increments) == {name: legacy[name] for name, _ in stored_daily.increments}
assert not {"messages", "response", "proxy_server_request", "requester_ip_address"}.intersection(accounting)
def test_projection_uses_frozen_savings_and_keeps_duplicate_entity_tags() -> None:
@ -316,3 +335,24 @@ def test_report_only_or_invalid_guardrail_prices_remain_untracked(cost_fields: M
model: Final = next(item for item in contributions if item.table == "LiteLLM_DailyModelUsage")
assert dict(guardrail.increments) == {"units": 2, "cost": 0.0, "untracked_units": 2}
assert dict(model.key)["task_type"] == "debugging"
@pytest.mark.parametrize("organization_id", [None, "", "charged-org"])
def test_daily_organization_uses_frozen_charged_identity(organization_id: str | None) -> None:
row: Final[Mapping[str, object]] = {
"startTime": _NOW,
"api_key": "key",
"model": "projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"organization_id": "stored-org",
"spend": 0.25,
"metadata": {"accounting_entities": {"organization_id": organization_id}},
}
daily: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyOrganizationSpend")
if organization_id is None:
assert daily == ()
else:
assert len(daily) == 1
assert dict(daily[0].key)["organization_id"] == organization_id
assert dict(daily[0].increments)["spend"] == 0.25

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

@ -753,17 +753,13 @@ class TestPartitionedSpendLogsPushGuard:
class _FakeCursor:
def __init__(self, present=True):
self._present = present
def fetchone(self):
return (1,) if self._present else None
return (1,)
class _FakePsycopgConn:
def __init__(self, executed, present=True):
def __init__(self, executed):
self._executed = executed
self._present = present
def __enter__(self):
return self
@ -773,7 +769,7 @@ class _FakePsycopgConn:
def execute(self, query, params):
self._executed.append((query, params))
return _FakeCursor(self._present)
return _FakeCursor()
class TestLensRenamePendingCheck:
@ -2377,42 +2373,3 @@ 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,15 +60,3 @@ 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;")