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": {}}