From 0faaeccfb65bf0838474fa2061e9e6d135e1f901 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Thu, 29 Oct 2026 13:48:44 -0700 Subject: [PATCH] test(proxy): pin the tag team upsert rule and drop the schema check changes The narrowed design no longer changes the daily tag spend unique key, and main's schema compatibility check already passes this schema, so its script and tests go back to main's version. Adds a test that a batch without a team never clears a stored team and one that the spend helper hands the virtual key's team to the tag rows, and drops a comment --- litellm/proxy/db/daily_spend_bulk_upsert.py | 2 - scripts/check_schema_compatibility.py | 40 ++------------ .../proxy/db/test_daily_spend_bulk_upsert.py | 7 +++ .../proxy/db/test_db_spend_update_writer.py | 39 ++++++++++++++ tests/unit/proxy/test_db_schema_changes.py | 52 +------------------ 5 files changed, 51 insertions(+), 89 deletions(-) diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index 74d4ee29029..2c20c521fb5 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -215,8 +215,6 @@ def build_bulk_upsert( if table.carries_request_id else "" ) - # Rows written by an older proxy, or by a request without a team's virtual key, carry "", so - # only a known team overwrites the stored one. team_id_update: Final = ( f', "team_id" = COALESCE(NULLIF(EXCLUDED."team_id", \'\'), {quoted_table}."team_id")' if table.carries_team_id diff --git a/scripts/check_schema_compatibility.py b/scripts/check_schema_compatibility.py index 541201fede8..898c65dff2c 100644 --- a/scripts/check_schema_compatibility.py +++ b/scripts/check_schema_compatibility.py @@ -12,47 +12,17 @@ from typing import Final from pydantic import TypeAdapter Schema = Mapping[str, Mapping[str, str]] -UniqueKeys = Mapping[str, frozenset[str]] MODELS: Final = re.compile(r"^[ \t]*model\s+(\w+)\s*\{(.*?)^[ \t]*\}", re.MULTILINE | re.DOTALL) -def _without_comment(tokens: list[str]) -> list[str]: - return tokens[: next((i for i, token in enumerate(tokens) if token.startswith("//")), len(tokens))] - - -def _is_field(part: list[str]) -> bool: - return len(part) >= 2 and not part[0].startswith("@@") - - def parse_model_fields(schema: str) -> Schema: def fields(body: str) -> Mapping[str, str]: - parts: Final = tuple(_without_comment(line.split()) for line in body.splitlines()) - return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if _is_field(part)}) + parts: Final = tuple(line.split() for line in body.splitlines()) + return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if len(part) >= 2}) return MappingProxyType({match[1]: fields(match[2]) for match in MODELS.finditer(schema)}) -def parse_unique_keys(schema: str) -> UniqueKeys: - def unique_keys(body: str) -> frozenset[str]: - parts: Final = tuple(_without_comment(line.split()) for line in body.splitlines()) - return frozenset(" ".join(part) for part in parts if part and part[0].startswith("@@unique")) - - return MappingProxyType({match[1]: unique_keys(match[2]) for match in MODELS.finditer(schema)}) - - -def check_unique_key_changes(old_keys: UniqueKeys, new_keys: UniqueKeys) -> tuple[str, ...]: - """An older proxy's upsert names its unique key in ON CONFLICT, so dropping or changing one breaks it.""" - - def changes() -> Iterator[str]: - for model, keys in old_keys.items(): - if model not in new_keys: - continue - for key in sorted(keys - new_keys[model]): - yield f"Breaking: Unique key {model}.{key} was removed or changed" - - return tuple(changes()) - - def _field_changes(model: str, field: str, old_type: str, new_fields: Mapping[str, str]) -> Iterator[str]: if field not in new_fields: yield f"Breaking: Field {model}.{field} was removed" @@ -93,10 +63,8 @@ def main() -> int: previous: Final = subprocess.run( ("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True ).stdout - current: Final = Path("schema.prisma").read_text() - changes: Final = ( - *check_breaking_changes(parse_model_fields(previous), parse_model_fields(current)), - *check_unique_key_changes(parse_unique_keys(previous), parse_unique_keys(current)), + changes: Final = check_breaking_changes( + parse_model_fields(previous), parse_model_fields(Path("schema.prisma").read_text()) ) if changes: sys.stderr.write("\n".join(changes) + "\n") diff --git a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py index 1f80e184516..af2d84c5169 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -134,6 +134,13 @@ def test_request_id_is_preserved_when_a_later_batch_carries_none(): assert None in params +def test_a_batch_without_a_team_never_clears_the_stored_team(): + sql, params = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(team_id=""),))) + + assert '"team_id" = COALESCE(NULLIF(EXCLUDED."team_id", \'\'), "LiteLLM_DailyTagSpend"."team_id")' in sql + assert "" in params + + def test_non_tag_tables_carry_no_request_id_column(): user_txn = {**tag_txn(), "user_id": "u-1"} del user_txn["tag"] diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 4dfbca7db2c..29ad4979a1b 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1344,6 +1344,45 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1} +@pytest.mark.asyncio +async def test_batch_database_updates_puts_the_virtual_key_team_on_tag_rows(): + db_writer: Final = DBSpendUpdateWriter() + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + await db_writer._batch_database_updates( + response_cost=0.01, + user_id="u1", + hashed_token="a" * 64, + team_id="team-a", + org_id=None, + end_user_id=None, + prisma_client=prisma, + litellm_proxy_budget_name=None, + payload=cast( + SpendLogsPayload, + { + "request_id": "req-1", + "request_tags": '["cost-center:cc-42"]', + "team_id": "team-a", + "user": "u1", + "startTime": "2026-10-29T00:00:00", + "api_key": "a" * 64, + "model": "gpt-4o", + "custom_llm_provider": "openai", + "model_group": "gpt-4o", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.01, + "metadata": '{"usage_object": {}}', + }, + ), + virtual_key_team_id="team-a", + ) + queued: Final = await db_writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + + assert [cast(DailyTagSpendTransaction, txn).get("team_id") for txn in queued.values()] == ["team-a"] + + @pytest.mark.asyncio async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted(): db_writer: Final = DBSpendUpdateWriter() diff --git a/tests/unit/proxy/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py index e6f9eb20dcc..3a5d0fb1a0a 100644 --- a/tests/unit/proxy/test_db_schema_changes.py +++ b/tests/unit/proxy/test_db_schema_changes.py @@ -2,12 +2,7 @@ from typing import Final import pytest -from scripts.check_schema_compatibility import ( - check_breaking_changes, - check_unique_key_changes, - parse_model_fields, - parse_unique_keys, -) +from scripts.check_schema_compatibility import check_breaking_changes, parse_model_fields ORIGINAL: Final = "model Item {\n id String @id\n label String?\n}\n" @@ -21,51 +16,6 @@ def test_parse_model_fields_retains_types_and_attributes() -> None: assert parse_model_fields(ORIGINAL) == {"Item": {"id": "String @id", "label": "String?"}} -def test_block_attributes_and_comments_are_not_fields() -> None: - schema: Final = ( - "model Item {\n id String @id\n label String?\n // free text\n" - ' @@unique([label, id], map: "item_label_id_key")\n @@index([label, id])\n @@index([id])\n}\n' - ) - assert parse_model_fields(schema) == {"Item": {"id": "String @id", "label": "String?"}} - - -@pytest.mark.parametrize( - "after", - [ - "model Item {\n id String @id\n label String?\n @@unique([label, id, tenant])\n}\n", - "model Item {\n id String @id\n label String?\n}\n", - ], - ids=["changed", "removed"], -) -def test_changing_or_removing_a_unique_key_is_breaking(after: str) -> None: - before: Final = "model Item {\n id String @id\n label String?\n @@unique([label, id])\n}\n" - assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == () - assert check_unique_key_changes(parse_unique_keys(before), parse_unique_keys(after)) == ( - "Breaking: Unique key Item.@@unique([label, id]) was removed or changed", - ) - - -def test_adding_an_index_or_rewording_a_comment_is_not_breaking() -> None: - before: Final = ( - 'model Item {\n id String @id\n label String @default("") // old note\n' - " @@unique([label, id])\n @@index([id])\n}\n" - ) - after: Final = ( - 'model Item {\n id String @id\n label String @default("") // reworded note\n' - " @@unique([label, id]) // same key\n @@index([id])\n @@index([label, id])\n}\n" - ) - assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == () - assert check_unique_key_changes(parse_unique_keys(before), parse_unique_keys(after)) == () - - -def test_a_narrowed_column_is_still_caught_next_to_block_attributes() -> None: - before: Final = "model Item {\n id String @id\n label String?\n @@index([label, id])\n}\n" - after: Final = "model Item {\n id String @id\n label String\n @@index([label, id])\n}\n" - assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == ( - "Breaking: Field Item.label changed from optional to required", - ) - - def test_parse_model_fields_keeps_json_defaults_and_empty_models() -> None: schema: Final = 'model Item {\n data Json @default("{}")\n}\nmodel Empty {\n}\n' assert parse_model_fields(schema) == {"Item": {"data": 'Json @default("{}")'}, "Empty": {}}