scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

refactor(proxy): name customer member budget columns budget_id and budget_override_id (#51490)

Browse source 
* refactor(proxy): rename customer member budget_id to budget_override_id

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* chore(ci): let the schema gate accept the unreleased customer member rename

* fix(proxy): refuse prisma db push while the customer member budget_id column is still unrenamed

* refactor(proxy): rename customer member budget_template_id to budget_id

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): refuse to baseline a database that still has the old customer member budget columns

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): suppress unknown member type on the optional proxy-extras budget rename guard

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-12 17:33:22 -08:00 • committed by GitHub
parent 24f753c589
commit c8036300bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 212 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
litellm/proxy-extras/litellm_proxy_extras/migrations/20261112120000_rename_customer_member_budget_columns/migration.sql Normal file
View file

@ -0,0 +1,36 @@
SET LOCAL lock_timeout = '10s';
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_CustomerMemberTable"')
AND attname = 'budget_template_id' AND NOT attisdropped
) THEN
RETURN;
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_CustomerMemberTable"')
AND attname = 'budget_override_id' AND NOT attisdropped
) THEN
ALTER TABLE "LiteLLM_CustomerMemberTable" RENAME COLUMN "budget_id" TO "budget_override_id";
END IF;
IF EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'LiteLLM_CustomerMemberTable_budget_id_fkey'
AND conrelid = to_regclass('"LiteLLM_CustomerMemberTable"')
) THEN
ALTER TABLE "LiteLLM_CustomerMemberTable" RENAME CONSTRAINT "LiteLLM_CustomerMemberTable_budget_id_fkey" TO "LiteLLM_CustomerMemberTable_budget_override_id_fkey";
END IF;
ALTER TABLE "LiteLLM_CustomerMemberTable" RENAME COLUMN "budget_template_id" TO "budget_id";
IF EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'LiteLLM_CustomerMemberTable_budget_template_id_fkey'
AND conrelid = to_regclass('"LiteLLM_CustomerMemberTable"')
) THEN
ALTER TABLE "LiteLLM_CustomerMemberTable" RENAME CONSTRAINT "LiteLLM_CustomerMemberTable_budget_template_id_fkey" TO "LiteLLM_CustomerMemberTable_budget_id_fkey";
END IF;
END $$;
ALTER INDEX IF EXISTS "LiteLLM_CustomerMemberTable_budget_template_id_idx" RENAME TO "LiteLLM_CustomerMemberTable_budget_id_idx";

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

@ -13,9 +13,9 @@ generator client {
model LiteLLM_BudgetTable {
customer_organization_defaults LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationDefaultBudget")
customer_organizations LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationBudget")
customer_member_templates LiteLLM_CustomerMemberTable[] @relation("CustomerMemberTemplate")
customer_teams LiteLLM_CustomerTeamTable[] @relation("CustomerTeamBudget")
customer_members LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudget")
customer_member_overrides LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudgetOverride")
budget_id String @id @default(uuid())
max_budget Float?
soft_budget Float?
@ -89,15 +89,15 @@ model LiteLLM_CustomerMemberTable {
parent_team LiteLLM_CustomerTeamTable @relation(fields: [parent_team_id], references: [customer_team_id])
budget_id String?
budget LiteLLM_BudgetTable? @relation("CustomerMemberBudget", fields: [budget_id], references: [budget_id], onDelete: Restrict)
budget_template_id String?
budget_template LiteLLM_BudgetTable? @relation("CustomerMemberTemplate", fields: [budget_template_id], references: [budget_id], onDelete: Restrict)
budget_override_id String?
budget_override LiteLLM_BudgetTable? @relation("CustomerMemberBudgetOverride", fields: [budget_override_id], references: [budget_id], onDelete: Restrict)
budget_override_expires_at DateTime?
created_at DateTime @default(now())
created_by String
updated_at DateTime @default(now()) @updatedAt
updated_by String
@@index([parent_team_id])
@@index([budget_template_id])
@@index([budget_id])
}
// Models on proxy

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

@ -643,37 +643,63 @@ class ProxyExtrasDBManager:
logger.warning(f"Failed to resolve migration {migration_name}: {e.stderr}")
@staticmethod
def raise_if_lens_rename_pending(connect: LensCheckConnect | None = None) -> None:
def _legacy_schema_object_exists(query: str, subject: str, connect: LensCheckConnect | None) -> bool:
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return
return False
try:
import psycopg
except ImportError as exc:
raise RuntimeError("Install psycopg to verify Lens data safety before prisma db push.") from exc
raise RuntimeError(f"Install psycopg to verify {subject} data safety before prisma db push.") from exc
open_connection: Final = connect if connect is not None else psycopg.connect
try:
with open_connection(
ProxyExtrasDBManager._strip_prisma_query_params(database_url), connect_timeout=10, autocommit=True
) as connection:
legacy: 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 IN ('LiteLLM_Engine', 'LiteLLM_EngineRun', 'LiteLLM_EngineWorker') "
"LIMIT 1",
(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",),
).fetchone()
return (
connection.execute(
query, (ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",)
).fetchone()
is not None
)
except psycopg.Error as exc:
raise RuntimeError(
"Cannot verify Lens data safety; refusing prisma db push. "
f"Cannot verify {subject} data safety; refusing prisma db push. "
f"The database check failed: {_redact_credentials(str(exc)).strip()}"
) from exc
if legacy is not None:
@staticmethod
def raise_if_lens_rename_pending(connect: LensCheckConnect | None = None) -> None:
if ProxyExtrasDBManager._legacy_schema_object_exists(
"SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid=c.relnamespace "
"WHERE n.nspname=%s AND c.relname IN ('LiteLLM_Engine', 'LiteLLM_EngineRun', 'LiteLLM_EngineWorker') "
"LIMIT 1",
"Lens",
connect,
):
raise RuntimeError(
"Legacy Lens tables exist. prisma db push would drop saved Lens data. "
"Apply the shipped 20261001100000_rename_lens migration to this database schema before retrying. "
"Deployments using migration history can upgrade without --use_prisma_db_push instead."
)
@staticmethod
def raise_if_customer_member_budget_rename_pending(connect: LensCheckConnect | None = None) -> None:
if ProxyExtrasDBManager._legacy_schema_object_exists(
"SELECT 1 FROM pg_attribute a JOIN pg_class c ON c.oid=a.attrelid "
"JOIN pg_namespace n ON n.oid=c.relnamespace "
"WHERE n.nspname=%s AND c.relname='LiteLLM_CustomerMemberTable' "
"AND a.attname='budget_template_id' AND NOT a.attisdropped LIMIT 1",
"customer member budget",
connect,
):
raise RuntimeError(
"LiteLLM_CustomerMemberTable still has the budget_template_id column. Syncing the schema would drop it "
"and treat every customer member budget override as the member's budget. Apply the shipped "
"20261112120000_rename_customer_member_budget_columns migration to this database schema "
"before retrying. Deployments with migration history upgrade without this step."
)
@staticmethod
def spend_logs_is_partitioned() -> bool:
"""True when the connected database's LiteLLM_SpendLogs is a
@ -1037,6 +1063,7 @@ class ProxyExtrasDBManager:
if not use_migrate:
ProxyExtrasDBManager.raise_if_lens_rename_pending()
ProxyExtrasDBManager.raise_if_customer_member_budget_rename_pending()
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
original_dir = os.getcwd()
@ -1623,6 +1650,7 @@ class ProxyExtrasDBManager:
logger.info(
"Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system"
)
ProxyExtrasDBManager.raise_if_customer_member_budget_rename_pending()
ProxyExtrasDBManager._create_baseline_migration(schema_path)
logger.info("Baseline migration created, resolving all migrations")
ProxyExtrasDBManager._resolve_all_migrations(migrations_dir, schema_path)
@ -1715,6 +1743,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_customer_member_budget_rename_pending()
prisma_toolchain.run_prisma(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),

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

@ -1281,6 +1281,7 @@ class PrismaManager:
from litellm_proxy_extras.utils import ProxyExtrasDBManager
ProxyExtrasDBManager.raise_if_lens_rename_pending()
ProxyExtrasDBManager.raise_if_customer_member_budget_rename_pending() # pyright: ignore[reportUnknownMemberType] # litellm-proxy-extras is an optional install
PrismaManager._raise_if_partitioned_spend_logs()
run_prisma(
[

44
litellm/proxy/management_helpers/customer_hierarchy.py
View file

@ -190,10 +190,10 @@ def effective_member_budget(
def member_response(row: PrismaRecord, now: datetime) -> CustomerMember:
values: Final = row.model_dump()
override: Final = related_budget(values.get("budget"))
template: Final = related_budget(values.get("budget_template"), now)
override: Final = related_budget(values.get("budget_override"))
template: Final = related_budget(values.get("budget"), now)
effective_override: Final = related_budget(
values.get("budget"), now, template.max_budget if template is not None else None
values.get("budget_override"), now, template.max_budget if template is not None else None
)
expiry: Final = _EXPIRY.validate_python(values.get("budget_override_expires_at"))
return CustomerMember.model_validate(
@ -223,7 +223,7 @@ def alias_patch(payload: OrganizationPatch | TeamPatch | MemberPatch) -> Mapping
_BUDGET_INCLUDE: Final = {"budget": True}
_MEMBER_INCLUDE: Final = {"budget": True, "budget_template": True}
_MEMBER_INCLUDE: Final = {"budget": True, "budget_override": True}
class CustomerHierarchyStore:
@ -437,7 +437,7 @@ class CustomerHierarchyStore:
return Failure(ErrorCode.TEMPLATE_NOT_IN_ORG)
if organization.default_budget_id == budget_id:
return Failure(ErrorCode.TEMPLATE_IN_USE)
if await self.members.count(where={"budget_template_id": budget_id}):
if await self.members.count(where={"budget_id": budget_id}):
return Failure(ErrorCode.TEMPLATE_IN_USE)
remaining: Final = {key: name for key, name in names.items() if key != budget_id}
await self.write_template_metadata(organization_id, remaining, organization.default_budget_id, actor)
@ -489,12 +489,12 @@ class CustomerHierarchyStore:
return await self.get_team(team_id)
async def create_member(self, payload: MemberCreate, actor: str) -> CustomerMember | Failure:
template_id: Final = await self.resolve_member_template(payload.parent_team_id, payload.budget_template_id)
template_id: Final = await self.resolve_member_template(payload.parent_team_id, payload.budget_id)
if isinstance(template_id, Failure):
return template_id
if template_id is None and payload.budget_override is None:
return Failure(ErrorCode.BUDGET_REQUIRED)
budget_id: Final = (
override_id: Final = (
await self.create_budget(payload.budget_override, actor) if payload.budget_override is not None else None
)
member_id: Final = payload.customer_member_id or str(uuid4())
@ -503,8 +503,8 @@ class CustomerHierarchyStore:
"customer_member_id": member_id,
"alias": payload.alias,
"parent_team_id": payload.parent_team_id,
"budget_template_id": template_id,
"budget_id": budget_id,
"budget_id": template_id,
"budget_override_id": override_id,
"budget_override_expires_at": payload.budget_override_expires_at,
"created_by": actor,
"updated_by": actor,
@ -513,18 +513,18 @@ class CustomerHierarchyStore:
return await self.get_member(member_id)
async def member_template_patch(self, member: CustomerMember, payload: MemberPatch) -> str | Failure | None:
if "budget_template_id" not in payload.model_fields_set:
return member.budget_template_id
if payload.budget_template_id is None:
if "budget_id" not in payload.model_fields_set:
return member.budget_id
if payload.budget_id is None:
return None
return await self.resolve_member_template(member.parent_team_id, payload.budget_template_id)
return await self.resolve_member_template(member.parent_team_id, payload.budget_id)
async def member_override_patch(self, member: CustomerMember, payload: MemberPatch, actor: str) -> str | None:
if "budget_override" not in payload.model_fields_set:
return member.budget_id
return member.budget_override_id
if payload.budget_override is None:
return None
return await self.update_budget(member.budget_id, payload.budget_override, actor)
return await self.update_budget(member.budget_override_id, payload.budget_override, actor)
async def update_member(self, member_id: str, payload: MemberPatch, actor: str) -> CustomerMember | Failure:
member: Final = await self.get_member(member_id)
@ -538,25 +538,25 @@ class CustomerHierarchyStore:
)
if selected is None and override is None:
return Failure(ErrorCode.BUDGET_REQUIRED)
budget_id: Final = await self.member_override_patch(member, payload, actor)
override_id: Final = await self.member_override_patch(member, payload, actor)
expiry: Final = (
payload.budget_override_expires_at
if "budget_override_expires_at" in payload.model_fields_set
else member.budget_override_expires_at
if budget_id is not None
if override_id is not None
else None
)
await self.members.update(
where={"customer_member_id": member_id},
data={
**alias_patch(payload),
"budget_id": budget_id,
"budget_template_id": selected,
"budget_id": selected,
"budget_override_id": override_id,
"budget_override_expires_at": expiry,
"updated_by": actor,
},
)
await self.delete_released_budget(member.budget_id, budget_id)
await self.delete_released_budget(member.budget_override_id, override_id)
return await self.get_member(member_id)
async def list_organizations(self, page: int, page_size: int, search: str = "") -> Page[Organization]:
@ -609,8 +609,8 @@ class CustomerHierarchyStore:
if isinstance(member, Failure):
return member
await self.members.delete(where={"customer_member_id": member_id})
if member.budget_id is not None:
await self.budgets.delete(where={"budget_id": member.budget_id})
if member.budget_override_id is not None:
await self.budgets.delete(where={"budget_id": member.budget_override_id})
return None
async def delete_team(self, team_id: str) -> None | Failure:

8
litellm/proxy/schema.prisma
View file

@ -13,9 +13,9 @@ generator client {
model LiteLLM_BudgetTable {
customer_organization_defaults LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationDefaultBudget")
customer_organizations LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationBudget")
customer_member_templates LiteLLM_CustomerMemberTable[] @relation("CustomerMemberTemplate")
customer_teams LiteLLM_CustomerTeamTable[] @relation("CustomerTeamBudget")
customer_members LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudget")
customer_member_overrides LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudgetOverride")
budget_id String @id @default(uuid())
max_budget Float?
soft_budget Float?
@ -89,15 +89,15 @@ model LiteLLM_CustomerMemberTable {
parent_team LiteLLM_CustomerTeamTable @relation(fields: [parent_team_id], references: [customer_team_id])
budget_id String?
budget LiteLLM_BudgetTable? @relation("CustomerMemberBudget", fields: [budget_id], references: [budget_id], onDelete: Restrict)
budget_template_id String?
budget_template LiteLLM_BudgetTable? @relation("CustomerMemberTemplate", fields: [budget_template_id], references: [budget_id], onDelete: Restrict)
budget_override_id String?
budget_override LiteLLM_BudgetTable? @relation("CustomerMemberBudgetOverride", fields: [budget_override_id], references: [budget_id], onDelete: Restrict)
budget_override_expires_at DateTime?
created_at DateTime @default(now())
created_by String
updated_at DateTime @default(now()) @updatedAt
updated_by String
@@index([parent_team_id])
@@index([budget_template_id])
@@index([budget_id])
}
// Models on proxy

8
litellm/schema.prisma
View file

@ -13,9 +13,9 @@ generator client {
model LiteLLM_BudgetTable {
customer_organization_defaults LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationDefaultBudget")
customer_organizations LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationBudget")
customer_member_templates LiteLLM_CustomerMemberTable[] @relation("CustomerMemberTemplate")
customer_teams LiteLLM_CustomerTeamTable[] @relation("CustomerTeamBudget")
customer_members LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudget")
customer_member_overrides LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudgetOverride")
budget_id String @id @default(uuid())
max_budget Float?
soft_budget Float?
@ -89,15 +89,15 @@ model LiteLLM_CustomerMemberTable {
parent_team LiteLLM_CustomerTeamTable @relation(fields: [parent_team_id], references: [customer_team_id])
budget_id String?
budget LiteLLM_BudgetTable? @relation("CustomerMemberBudget", fields: [budget_id], references: [budget_id], onDelete: Restrict)
budget_template_id String?
budget_template LiteLLM_BudgetTable? @relation("CustomerMemberTemplate", fields: [budget_template_id], references: [budget_id], onDelete: Restrict)
budget_override_id String?
budget_override LiteLLM_BudgetTable? @relation("CustomerMemberBudgetOverride", fields: [budget_override_id], references: [budget_id], onDelete: Restrict)
budget_override_expires_at DateTime?
created_at DateTime @default(now())
created_by String
updated_at DateTime @default(now()) @updatedAt
updated_by String
@@index([parent_team_id])
@@index([budget_template_id])
@@index([budget_id])
}
// Models on proxy

13
litellm/scripts/check_schema_compatibility.py
View file

@ -14,6 +14,13 @@ from pydantic import TypeAdapter
Schema = Mapping[str, Mapping[str, str]]
MODELS: Final = re.compile(r"^[ \t]*model\s+(\w+)\s*\{(.*?)^[ \t]*\}", re.MULTILINE | re.DOTALL)
DROPPED_UNUSED_MODELS: Final = frozenset({"LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"})
RENAMED_UNRELEASED_FIELDS: Final = frozenset(
{
("LiteLLM_CustomerMemberTable", "budget_template_id"),
("LiteLLM_CustomerMemberTable", "budget_template"),
("LiteLLM_BudgetTable", "customer_member_templates"),
}
)
def parse_model_fields(schema: str) -> Schema:
@ -26,6 +33,8 @@ def parse_model_fields(schema: str) -> Schema:
def _field_changes(model: str, field: str, old_type: str, new_fields: Mapping[str, str]) -> Iterator[str]:
if field not in new_fields:
if (model, field) in RENAMED_UNRELEASED_FIELDS:
return
yield f"Breaking: Field {model}.{field} was removed"
return
old_base: Final = old_type.split()[0]
@ -73,9 +82,7 @@ def main() -> int:
previous: Final = (
canonical.stdout
if canonical.returncode == 0
else subprocess.run(
("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True
).stdout
else subprocess.run(("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True).stdout
)
changes: Final = check_breaking_changes(
parse_model_fields(previous), parse_model_fields(Path("litellm/schema.prisma").read_text())

4
litellm/types/proxy/management_endpoints/customer_hierarchy.py
View file

@ -67,8 +67,8 @@ class MemberPatch(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
alias: str | None = None
budget_id: str | None = None
budget_override: Budget | None = None
budget_template_id: str | None = None
budget_override_expires_at: AwareDatetime | None = None
@ -109,7 +109,7 @@ class CustomerMember(AuditFields):
customer_member_id: str
parent_team_id: str
budget_id: str | None = None
budget_template_id: str | None = None
budget_override_id: str | None = None
budget_override: Budget | None = None
budget_override_expires_at: datetime | None = None
effective_budget: Budget

12
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -33137,10 +33137,10 @@ export interface components {
budget_override?: components["schemas"]["Budget"] | null;
/** Budget Override Expires At */
budget_override_expires_at?: string | null;
/** Budget Override Id */
budget_override_id?: string | null;
/** Budget Reset At */
budget_reset_at?: string | null;
/** Budget Template Id */
budget_template_id?: string | null;
/**
* Created At
* Format: date-time
@ -40405,11 +40405,11 @@ export interface components {
MemberCreate: {
/** Alias */
alias?: string | null;
/** Budget Id */
budget_id?: string | null;
budget_override?: components["schemas"]["Budget"] | null;
/** Budget Override Expires At */
budget_override_expires_at?: string | null;
/** Budget Template Id */
budget_template_id?: string | null;
/** Customer Member Id */
customer_member_id?: string | null;
/** Parent Team Id */
@ -40419,11 +40419,11 @@ export interface components {
MemberPatch: {
/** Alias */
alias?: string | null;
/** Budget Id */
budget_id?: string | null;
budget_override?: components["schemas"]["Budget"] | null;
/** Budget Override Expires At */
budget_override_expires_at?: string | null;
/** Budget Template Id */
budget_template_id?: string | null;
};
/** MemoryCreateRequest */
MemoryCreateRequest: {

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

@ -818,6 +818,34 @@ class TestLensRenamePendingCheck:
assert executed[0][1] == ("public",)
class TestCustomerMemberBudgetRenamePendingCheck:
_DATABASE_URL: Final = "postgresql://litellm:hunter2@localhost:5432/litellm?schema=tenant_a"
def test_a_leftover_budget_template_id_column_refuses_the_push(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DATABASE_URL", self._DATABASE_URL)
executed: Final[list[tuple[str, tuple[str]]]] = []
def legacy_column_present(conninfo: str, *, connect_timeout: int, autocommit: bool) -> _FakePsycopgConn:
return _FakePsycopgConn(executed)
with pytest.raises(RuntimeError) as err:
ProxyExtrasDBManager.raise_if_customer_member_budget_rename_pending(connect=legacy_column_present)
assert "20261112120000_rename_customer_member_budget_columns" in str(err.value)
assert executed[0][1] == ("tenant_a",)
def test_an_unreachable_database_refuses_the_push(self, monkeypatch: pytest.MonkeyPatch) -> None:
import psycopg
monkeypatch.setenv("DATABASE_URL", self._DATABASE_URL)
def refuse(conninfo: str, *, connect_timeout: int, autocommit: bool) -> NoReturn:
raise psycopg.OperationalError("FATAL: sorry, too many clients already")
with pytest.raises(RuntimeError) as err:
ProxyExtrasDBManager.raise_if_customer_member_budget_rename_pending(connect=refuse)
assert "FATAL: sorry, too many clients already" in str(err.value)
class TestSpendLogsPartitionDetectionSchemaScope:
"""A same-named LiteLLM_SpendLogs in another schema must not trip the
detector: the catalog lookup has to be scoped to Prisma's target schema."""
@ -1886,9 +1914,26 @@ class TestV1MigrationFailuresLogAtError:
deploy_timeout: bool = False,
diff_stderr: Optional[str] = None,
database_url: Optional[str] = None,
member_budget_rename_pending: bool = False,
) -> tuple[bool, list[list[str]], list[int]]:
import litellm_proxy_extras.utils as utils_module
class _MemberColumnCheck:
def __enter__(self) -> "_MemberColumnCheck":
return self
def __exit__(self, *args: object) -> None:
return None
def execute(self, query: str, params: tuple[str]) -> "_MemberColumnCheck":
return self
def fetchone(self) -> tuple[int] | None:
return (1,) if member_budget_rename_pending else None
monkeypatch.setattr(
"psycopg.connect", lambda conninfo, *, connect_timeout, autocommit: _MemberColumnCheck()
)
calls: Final[list[list[str]]] = []
killed_pids: Final[list[int]] = []
@ -2096,6 +2141,19 @@ class TestV1MigrationFailuresLogAtError:
assert tuple(calls[0][-2:]) == ("migrate", "deploy")
assert self._error_messages(caplog) == []
def test_a_baseline_refuses_while_the_member_budget_rename_is_pending(
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
with pytest.raises(RuntimeError) as err:
self._run_v1_migrations(
monkeypatch,
tmp_path,
deploy_stderr=_P3005_STDERR,
database_url="postgresql://user:pass@db:5432/litellm",
member_budget_rename_pending=True,
)
assert "20261112120000_rename_customer_member_budget_columns" in str(err.value)
def test_a_failed_baseline_recovery_logs_its_stderr_at_error(
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture
) -> None:

22
tests/unit/proxy/management_endpoints/management_v1/test_customers.py
View file

@ -113,14 +113,14 @@ def test_template_edits_propagate_and_only_live_override_fields_win() -> None:
},
)
assert default.status_code == 201, default.text
assert default.json()["budget_id"] is None
assert default.json()["budget_template_id"] == default_id
assert default.json()["budget_override_id"] is None
assert default.json()["budget_id"] == default_id
override: Final = http.post(
"/management/v1/customer_members",
json={
"customer_member_id": "override",
"parent_team_id": "team",
"budget_template_id": other_id,
"budget_id": other_id,
"budget_override": {"max_budget": 3, "allowed_models": []},
},
)
@ -201,15 +201,15 @@ def test_override_can_be_reexpired_cleared_and_template_switched() -> None:
switched: Final = http.patch(
"/management/v1/customer_members/member",
json={
"budget_template_id": org["default_budgets"][1]["budget_id"],
"budget_id": org["default_budgets"][1]["budget_id"],
"budget_override": None,
},
)
assert switched.status_code == 200, switched.text
assert switched.json()["budget_id"] is None
assert switched.json()["budget_override_id"] is None
assert switched.json()["budget_override"] is None
assert switched.json()["effective_budget"]["max_budget"] == 20
cleared: Final = http.patch("/management/v1/customer_members/member", json={"budget_template_id": None})
cleared: Final = http.patch("/management/v1/customer_members/member", json={"budget_id": None})
assert cleared.status_code == 422, cleared.text
after: Final = http.get("/management/v1/customer_members/member")
assert after.json() == switched.json()
@ -268,7 +268,7 @@ def test_template_deletion_guards_names_and_references() -> None:
json={
"customer_member_id": "member",
"parent_team_id": "team",
"budget_template_id": selected_id,
"budget_id": selected_id,
},
)
assert member.status_code == 201, member.text
@ -312,7 +312,7 @@ def test_missing_template_requires_custom_budget_and_wrong_org_is_rejected() ->
"/management/v1/customer_members",
json={
"parent_team_id": "empty-team",
"budget_template_id": org["default_budget_id"],
"budget_id": org["default_budget_id"],
"budget_override": {"max_budget": 2},
},
)
@ -321,7 +321,7 @@ def test_missing_template_requires_custom_budget_and_wrong_org_is_rejected() ->
"/management/v1/customer_members", json={"parent_team_id": "empty-team", "budget_override": {"max_budget": 2}}
)
assert custom.status_code == 201, custom.text
assert custom.json()["budget_template_id"] is None
assert custom.json()["budget_id"] is None
assert custom.json()["effective_budget"]["max_budget"] == 2
@ -487,7 +487,7 @@ def test_partial_member_override_updates_preserve_other_fields_and_never_change_
},
)
assert patched.status_code == 200, patched.text
assert patched.json()["budget_id"] != organization["default_budget_id"]
assert patched.json()["budget_override_id"] != organization["default_budget_id"]
assert patched.json()["effective_budget"]["max_budget"] == 5
inherited: Final = http.patch(
"/management/v1/customer_members/member",
@ -496,7 +496,7 @@ def test_partial_member_override_updates_preserve_other_fields_and_never_change_
},
)
assert inherited.status_code == 200, inherited.text
assert inherited.json()["budget_id"] == patched.json()["budget_id"]
assert inherited.json()["budget_override_id"] == patched.json()["budget_override_id"]
assert inherited.json()["effective_budget"]["max_budget"] == 100
assert inherited.json()["effective_budget"]["soft_budget"] == 4
assert inherited.json()["effective_budget"]["tpm_limit"] == 2

14
tests/unit/proxy/test_db_schema_changes.py
View file

@ -60,3 +60,17 @@ def test_only_the_dropped_accounting_models_may_be_removed() -> None:
assert check_breaking_changes(
previous, parse_model_fields(ORIGINAL + "model LiteLLM_AccountingDelivery {\n event_id String\n}\n")
) == ("Breaking: Field LiteLLM_AccountingDelivery.note was removed",)
def test_only_the_renamed_customer_member_fields_may_be_removed() -> None:
member: Final = (
"model LiteLLM_CustomerMemberTable {\n alias String?\n budget_template_id String?\n budget_template Item?\n}\n"
)
previous: Final = parse_model_fields(ORIGINAL + member)
renamed: Final = "model LiteLLM_CustomerMemberTable {\n alias String?\n budget_id String?\n}\n"
without_alias: Final = "model LiteLLM_CustomerMemberTable {\n budget_id String?\n}\n"
assert check_breaking_changes(previous, parse_model_fields(ORIGINAL + renamed)) == ()
assert check_breaking_changes(previous, parse_model_fields(ORIGINAL + without_alias)) == (
"Breaking: Field LiteLLM_CustomerMemberTable.alias was removed",
)