diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql new file mode 100644 index 00000000000..0b73df2d05b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql @@ -0,0 +1,10 @@ +SET LOCAL lock_timeout = '10s'; + +-- Records the team that sent each request on its daily tag spend rows, so one tag used by +-- several teams lands in one row per team. "" means the request had no team, and every row +-- written before this migration keeps "" (attribution is not retroactive). +-- +-- NOT NULL DEFAULT '' rather than nullable: the daily spend upsert normalizes every conflict +-- column to '' because a NULL never matches itself in a unique index. On PostgreSQL 11+ adding +-- a column with a constant default is a metadata-only change, so this does not rewrite the table. +ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "team_id" TEXT NOT NULL DEFAULT ''; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql new file mode 100644 index 00000000000..1816ff63ce6 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql @@ -0,0 +1,14 @@ +-- CreateIndex (CONCURRENTLY) +-- +-- The new conflict target for the daily tag spend upsert: the old unique key plus team_id. +-- Built before the old unique index is dropped (next migrations), so the table is never +-- without a unique index the writer can arbitrate on. +-- +-- Disclaimer: +-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - If the build is interrupted, Postgres may leave an INVALID index that must be dropped and +-- recreated. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key" ON "LiteLLM_DailyTagSpend"("tag", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql new file mode 100644 index 00000000000..474d5311305 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql @@ -0,0 +1,12 @@ +-- CreateIndex (CONCURRENTLY) +-- +-- Serves team-scoped tag reads: WHERE team_id = ANY(...) AND date BETWEEN ... grouped by tag. +-- +-- Disclaimer: +-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - If the build is interrupted, Postgres may leave an INVALID index that must be dropped and +-- recreated. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_team_id_date_tag_idx" ON "LiteLLM_DailyTagSpend"("team_id", "date", "tag"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql new file mode 100644 index 00000000000..ebc4f7ac457 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql @@ -0,0 +1,18 @@ +-- DropIndex (CONCURRENTLY) +-- +-- Drops the unique key without team_id. With it in place, two teams sending the same tag on +-- the same day for the same key and model would collide on it, so it has to go once the +-- team-aware unique index exists. +-- +-- Upgrade note: an old-version writer still targets this index in its ON CONFLICT clause and +-- fails every flush once it is gone. Drain and stop old instances before running migrations. +-- An old instance booting against the migrated database also logs "prisma schema out of sync" +-- with SQL that drops team_id and both new indexes. Do not run it: it deletes team attribution +-- and its CREATE UNIQUE INDEX fails once two teams share a key. +-- +-- Disclaimer: +-- - DROP INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 83a0e6f9e05..c84b9f810fe 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2b01e02bb06..ddf1760a554 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5223,6 +5223,7 @@ class DailyEndUserSpendTransaction(BaseDailySpendTransaction): class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: str | None tag: str + team_id: ReadOnly[NotRequired[str]] class DailyAgentSpendTransaction(BaseDailySpendTransaction): diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index f7a2be654c8..5e817410d2b 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -58,6 +58,7 @@ class DailyBaselineAttribution(LiteLLMBaseModel): model_group: str | None = None endpoint: str | None = None mcp_namespaced_tool_name: str | None = None + team_id: str = "" targets: tuple[DailyBaselineTarget, ...] = () def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow: @@ -71,6 +72,7 @@ class DailyBaselineAttribution(LiteLLMBaseModel): "model_group": self.model_group, "endpoint": self.endpoint, "mcp_namespaced_tool_name": self.mcp_namespaced_tool_name, + **({"team_id": self.team_id} if table.team_scoped else {}), table.entity_id_column: target.entity_id, "request_id": request_id, "autorouter_savings_spend": savings_delta, diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index eb130a5196f..dcc2b738a51 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -32,6 +32,7 @@ class DailySpendTable: name: str entity_id_column: str carries_request_id: bool = False + team_scoped: bool = False DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( @@ -41,7 +42,9 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), - "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), + "tag": DailySpendTable( + name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True, team_scoped=True + ), } ) @@ -124,9 +127,14 @@ def _as_float(value: object) -> float: return float(value) if isinstance(value, (int, float)) else 0.0 +def _conflict_columns(table: DailySpendTable) -> tuple[str, ...]: + """The table's unique constraint, in constraint order.""" + return (table.entity_id_column, *(("team_id",) if table.team_scoped else ()), *_KEY_COLUMNS) + + def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) + return tuple(_as_text(transaction.get(column)) for column in _conflict_columns(table)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -174,8 +182,7 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - table.entity_id_column, - *_KEY_COLUMNS, + *_conflict_columns(table), "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, @@ -213,7 +220,7 @@ def build_bulk_upsert( sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted(_conflict_columns(table))}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 8ac2597c99b..7870ee15d2f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1019,6 +1019,7 @@ class DBSpendUpdateWriter: model_group=base.get("model_group"), endpoint=base.get("endpoint"), mcp_namespaced_tool_name=base.get("mcp_namespaced_tool_name"), + team_id=payload.get("team_id") or "", # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults targets=tuple( DailyBaselineTarget(entity=entity, entity_id=identity) for entity, values in bases @@ -3075,13 +3076,25 @@ class DBSpendUpdateWriter: return request_tags: Final = daily_spend_entity_ids(payload, "tag") + # .get, not indexing: the caller swallows exceptions, so a KeyError would silently drop tag spend + team_id: Final = payload.get("team_id") or "" # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults for tag in request_tags: if tag is None: continue endpoint_str = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" + daily_transaction_key = json.dumps( + [ + tag, + team_id, + base_daily_transaction["date"], + payload["api_key"], + payload["model"], + payload["custom_llm_provider"], + endpoint_str, + ] + ) daily_transaction = DailyTagSpendTransaction( - tag=tag, **base_daily_transaction, request_id=payload["request_id"] + tag=tag, team_id=team_id, **base_daily_transaction, request_id=payload["request_id"] ) await self.daily_tag_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 83a0e6f9e05..c84b9f810fe 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) diff --git a/schema.prisma b/schema.prisma index 83a0e6f9e05..c84b9f810fe 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) diff --git a/scripts/check_schema_compatibility.py b/scripts/check_schema_compatibility.py index 898c65dff2c..5e98db3c7e2 100644 --- a/scripts/check_schema_compatibility.py +++ b/scripts/check_schema_compatibility.py @@ -15,10 +15,14 @@ Schema = Mapping[str, Mapping[str, str]] MODELS: Final = re.compile(r"^[ \t]*model\s+(\w+)\s*\{(.*?)^[ \t]*\}", re.MULTILINE | re.DOTALL) +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(line.split() for line in body.splitlines()) - return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if len(part) >= 2}) + return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if _is_field(part)}) return MappingProxyType({match[1]: fields(match[2]) for match in MODELS.finditer(schema)}) diff --git a/tests/integration/spend/test_team_tag_spend_attribution.py b/tests/integration/spend/test_team_tag_spend_attribution.py new file mode 100644 index 00000000000..d82394e143b --- /dev/null +++ b/tests/integration/spend/test_team_tag_spend_attribution.py @@ -0,0 +1,45 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows + +_SPEND_PER_CALL: Final = 20 * 0.001 + 20 * 0.002 + + +def test_one_key_moved_between_teams_records_its_tag_once_per_team(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + tag: Final = f"cost-center:{uuid.uuid4().hex}" + teams: Final = tuple(scenario.team(models=[model]) for _ in range(2)) + calls_per_team: Final = {teams[0]: 1, teams[1]: 2} + key: Final = scenario.key(team_id=teams[0], models=[model]) + for team, calls in calls_per_team.items(): + gateway.post("/key/update", {"key": key, "team_id": team}) + for _ in range(calls): + reply = gateway.post( + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"tagged {uuid.uuid4().hex}"}], + "metadata": {"tags": [tag]}, + }, + key=key, + ) + assert reply["usage"]["total_tokens"] == 40, reply + + rows: Final = eventually( + lambda: read_rows( + 'SELECT team_id, spend, api_requests FROM "LiteLLM_DailyTagSpend" WHERE tag = %s ORDER BY team_id', + (tag,), + ), + lambda values: sum(int(row["api_requests"]) for row in values) == 3, + seconds=70, + ) + + assert sorted((row["team_id"], int(row["api_requests"])) for row in rows) == sorted(calls_per_team.items()), ( + rows + ) + for row in rows: + assert float(row["spend"]) == pytest.approx(calls_per_team[row["team_id"]] * _SPEND_PER_CALL), rows 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 02bdf3429ea..ad889af1873 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 @@ -435,7 +435,7 @@ class TestMigrationSQLIdempotency: for migration_name, sql in _guarded_migrations(all_migrations): for line_num, line in _statements(sql): if re.search(r"DROP\s+INDEX\s+", line, re.IGNORECASE) and not re.search( - r"DROP\s+INDEX\s+IF\s+EXISTS", line, re.IGNORECASE + r"DROP\s+INDEX\s+(?:CONCURRENTLY\s+)?IF\s+EXISTS", line, re.IGNORECASE ): violations.append(f" {migration_name}:{line_num}: {line.strip()}") assert ( 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 5bce9137666..61a7d7dac7f 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -3,11 +3,13 @@ import re from collections.abc import AsyncIterator from contextlib import AbstractAsyncContextManager, asynccontextmanager +from typing import Final import pytest from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, + DailySpendEntity, build_bulk_upsert, conflict_key, merge_by_conflict_key, @@ -88,10 +90,8 @@ def test_one_statement_carries_every_row_in_the_batch(): assert sql.count("INSERT INTO") == 1 assert len(re.findall(r"ON CONFLICT", sql)) == 1 - # 25 bound columns per row plus the inlined updated_at, so the row count is what - # separates one multi-row statement from a hundred single-row ones. - assert len(params) == 100 * 25 - assert "$2500::text" in sql + assert len(params) == 100 * 26 + assert "$2600::text" in sql assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1 @@ -101,7 +101,7 @@ def test_conflict_target_is_the_full_unique_constraint(): conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) assert conflict_target is not None assert conflict_target.group(1) == ( - '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' + '"tag", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' ) @@ -144,6 +144,41 @@ def test_non_tag_tables_carry_no_request_id_column(): assert '"user_id"' in sql +def test_one_tag_from_two_teams_stays_two_rows() -> None: + merged: Final = merge_by_conflict_key( + TAG_TABLE, (tag_txn(team_id="team-b", spend=2.0), tag_txn(team_id="team-a", spend=1.0)) + ) + + assert [(txn["team_id"], txn["spend"]) for _, txn in merged] == [("team-a", 1.0), ("team-b", 2.0)], merged + + +def test_an_entry_without_team_id_folds_into_the_unattributed_row() -> None: + legacy: Final = tag_txn(spend=0.25) + unattributed: Final = tag_txn(team_id="", spend=0.75) + + merged: Final = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed)) + + assert len(merged) == 1, merged + key, folded = merged[0] + assert key[:2] == ("team-a", ""), key + assert folded["spend"] == pytest.approx(1.0) + + +@pytest.mark.parametrize("entity", ["user", "team", "org", "end_user", "agent"]) +def test_only_the_tag_table_is_keyed_by_team(entity: DailySpendEntity) -> None: + table: Final = DAILY_SPEND_TABLES[entity] + row: Final = {**tag_txn(), table.entity_id_column: "id-1"} + + sql, _ = build_bulk_upsert(table, merge_by_conflict_key(table, (row,))) + + conflict_target: Final = re.search(r"ON CONFLICT \(([^)]*)\)", sql) + assert conflict_target is not None + assert conflict_target.group(1) == ( + f'"{table.entity_id_column}", "date", "api_key", "model", "custom_llm_provider", ' + '"mcp_namespaced_tool_name", "endpoint"' + ) + + class _RecordingDb: def __init__(self) -> None: self.statements: list[tuple[str, tuple[object, ...]]] = [] 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 a42b6404278..531e7fb787c 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -20,7 +20,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS -from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem +from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendLogsPayload, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, _TEAM_MEMBER_SPEND_SQL, @@ -1493,6 +1493,67 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i ) +def _tagged_spend_log(tag: str, team_id: str | None, spend: float) -> SpendLogsPayload: + return cast( + SpendLogsPayload, + { + "request_id": f"req-{tag}-{team_id}", + "request_tags": json.dumps([tag]), + "team_id": team_id, + "user": "user-1", + "startTime": "2026-10-17T00:00:00", + "api_key": "key-1", + "model": "gpt-4o", + "custom_llm_provider": "openai", + "model_group": "gpt-4o", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": spend, + "metadata": '{"usage_object": {}}', + }, + ) + + +@pytest.mark.parametrize( + ("requests", "expected"), + [ + ( + ( + ("cost-center:cc-42", "team-a", 1.0), + ("cost-center:cc-42", "team-b", 2.0), + ("cost-center:cc-42", None, 0.5), + ), + [ + ("cost-center:cc-42", "", 0.5), + ("cost-center:cc-42", "team-a", 1.0), + ("cost-center:cc-42", "team-b", 2.0), + ], + ), + ( + (("cost_center", "west", 1.0), ("cost", "center_west", 2.0)), + [("cost", "center_west", 2.0), ("cost_center", "west", 1.0)], + ), + ], + ids=["one-tag-two-teams", "underscore-joined-names-collide"], +) +async def test_each_tag_and_team_pair_is_queued_as_its_own_entry( + requests: tuple[tuple[str, str | None, float], ...], + expected: list[tuple[str, str, float]], +) -> None: + writer: Final = DBSpendUpdateWriter() + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + + for tag, team_id, spend in requests: + await writer.add_spend_log_transaction_to_daily_tag_transaction( + payload=_tagged_spend_log(tag, team_id, spend), prisma_client=prisma + ) + + queued: Final = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + + assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == expected, queued + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id_and_queues_update(): """ diff --git a/tests/unit/proxy/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py index 3a5d0fb1a0a..49dfaec1f16 100644 --- a/tests/unit/proxy/test_db_schema_changes.py +++ b/tests/unit/proxy/test_db_schema_changes.py @@ -16,6 +16,28 @@ 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?"}} + + +def test_changing_a_composite_unique_key_is_not_a_field_change() -> None: + before: Final = "model Item {\n id String @id\n label String?\n @@unique([label, id])\n}\n" + after: Final = "model Item {\n id String @id\n label String?\n @@unique([label, id, tenant])\n}\n" + assert check_breaking_changes(parse_model_fields(before), parse_model_fields(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": {}}