diff --git a/litellm/proxy-extras/litellm_proxy_extras/migrations/20261112120000_rename_customer_member_budget_columns/migration.sql b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261112120000_rename_customer_member_budget_columns/migration.sql new file mode 100644 index 00000000000..9d93f5530c1 --- /dev/null +++ b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261112120000_rename_customer_member_budget_columns/migration.sql @@ -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"; diff --git a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma index bac27daa5de..3e7e18096da 100644 --- a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma @@ -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 diff --git a/litellm/proxy-extras/litellm_proxy_extras/utils.py b/litellm/proxy-extras/litellm_proxy_extras/utils.py index f8625088d8e..2b988df14f6 100644 --- a/litellm/proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm/proxy-extras/litellm_proxy_extras/utils.py @@ -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(), diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index e0be2e1e5ef..1b9851cf082 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -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( [ diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py index 3ec37300127..3ff6c7846e2 100644 --- a/litellm/proxy/management_helpers/customer_hierarchy.py +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -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: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index bac27daa5de..3e7e18096da 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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 diff --git a/litellm/schema.prisma b/litellm/schema.prisma index bac27daa5de..3e7e18096da 100644 --- a/litellm/schema.prisma +++ b/litellm/schema.prisma @@ -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 diff --git a/litellm/scripts/check_schema_compatibility.py b/litellm/scripts/check_schema_compatibility.py index 5d7ee790cb2..b529e839caf 100644 --- a/litellm/scripts/check_schema_compatibility.py +++ b/litellm/scripts/check_schema_compatibility.py @@ -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()) diff --git a/litellm/types/proxy/management_endpoints/customer_hierarchy.py b/litellm/types/proxy/management_endpoints/customer_hierarchy.py index aca4c1626df..db6cc101b2f 100644 --- a/litellm/types/proxy/management_endpoints/customer_hierarchy.py +++ b/litellm/types/proxy/management_endpoints/customer_hierarchy.py @@ -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 diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index 64e6fabfdb9..e305581093a 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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: { 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 186a38c4327..67f51055de5 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 @@ -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: diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py index 98bba1ac043..f9763ad6855 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -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 diff --git a/tests/unit/proxy/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py index c2e72c86329..5ed8c4723e9 100644 --- a/tests/unit/proxy/test_db_schema_changes.py +++ b/tests/unit/proxy/test_db_schema_changes.py @@ -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", + )