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

@@ -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

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")

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"]

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()

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