From 7bc136d8212595af3305823ddc711738aa8e3b6b Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 3 Oct 2026 11:20:48 -0700 Subject: [PATCH] Revert "fix(usage): match legacy NULL model_group through COALESCE instead of a backfill" This reverts commit 29607f8794487812d1251ccd4b6dd2f5cc9f0e00. Prisma cannot declare the COALESCE expression index, so every database seeded by `prisma db push` outside the proxy's own startup (the proxy-behavior CI job, the e2e UI script) had no unique index on the daily spend tables and the writer's ON CONFLICT failed. The legacy NULL model_group duplicate stays a documented cutover caveat. Co-Authored-By: Claude Fable 5.1 --- .../migration.sql | 12 +++--- .../litellm_proxy_extras/schema.prisma | 6 +++ .../litellm_proxy_extras/utils.py | 23 ----------- litellm-proxy-extras/migration_runbook.md | 2 +- litellm/proxy/db/daily_spend_bulk_upsert.py | 13 +------ litellm/proxy/schema.prisma | 6 +++ schema.prisma | 6 +++ .../test_litellm_proxy_extras_utils.py | 39 ------------------- .../proxy/db/test_daily_spend_bulk_upsert.py | 7 ++-- 9 files changed, 29 insertions(+), 85 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql index 3f52ce49dc0..bc803c0a096 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql @@ -1,16 +1,16 @@ BEGIN; -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_group_key" ON "LiteLLM_DailyUserSpend" ("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", '')); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_group_key" ON "LiteLLM_DailyUserSpend" ("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_group_key" ON "LiteLLM_DailyTeamSpend" ("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", '')); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_group_key" ON "LiteLLM_DailyTeamSpend" ("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_group_key" ON "LiteLLM_DailyOrganizationSpend" ("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", '')); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_group_key" ON "LiteLLM_DailyOrganizationSpend" ("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_group_key" ON "LiteLLM_DailyEndUserSpend" ("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", '')); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_group_key" ON "LiteLLM_DailyEndUserSpend" ("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_group_key" ON "LiteLLM_DailyAgentSpend" ("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", '')); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_group_key" ON "LiteLLM_DailyAgentSpend" ("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_group_key" ON "LiteLLM_DailyTagSpend" ("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", '')); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_group_key" ON "LiteLLM_DailyTagSpend" ("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f1a380356ce..cbd31d98814 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -873,6 +873,7 @@ model LiteLLM_DailyUserSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key") @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -941,6 +942,7 @@ model LiteLLM_DailyOrganizationSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key") @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -977,6 +979,7 @@ model LiteLLM_DailyEndUserSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key") @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -1013,6 +1016,7 @@ model LiteLLM_DailyAgentSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key") @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -1051,6 +1055,7 @@ model LiteLLM_DailyTeamSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key") @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -1089,6 +1094,7 @@ 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, model_group], map: "LiteLLM_DailyTagSpend_group_key") @@index([date]) @@index([tag, date]) @@index([api_key]) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index abacaf7ca39..245244250ee 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -202,9 +202,6 @@ def _max_migration_timestamp(names) -> int: return max(_migration_timestamp(n) for n in names) -DAILY_SPEND_IDENTITY_MIGRATION: Final = "20260922000000_daily_spend_model_group_identity" - - def _get_prisma_command() -> str: """Get the Prisma command to use, bypassing Python wrapper in offline mode.""" if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): @@ -1030,7 +1027,6 @@ class ProxyExtrasDBManager: stdout=None, stderr=None, ) - ProxyExtrasDBManager._apply_daily_spend_identity(migrations_dir) return True except ( subprocess.CalledProcessError, @@ -1355,24 +1351,6 @@ class ProxyExtrasDBManager: thread.start() return thread - @staticmethod - def _apply_daily_spend_identity(prisma_dir: str) -> None: - """`prisma db push` builds only what schema.prisma declares. The daily spend identity is a - COALESCE expression index Prisma cannot declare, so the migration owning it runs after the push.""" - prisma_toolchain.run_prisma( - [ - _get_prisma_command(), - "db", - "execute", - "--file", - f"{prisma_dir}/migrations/{DAILY_SPEND_IDENTITY_MIGRATION}/migration.sql", - "--schema", - f"{prisma_dir}/schema.prisma", - ], - timeout=prisma_command_timeout(), - env=_get_prisma_env(), - ) - @staticmethod def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool: if use_v2_resolver: @@ -1606,7 +1584,6 @@ class ProxyExtrasDBManager: stderr=None, env=_get_prisma_env(), ) - ProxyExtrasDBManager._apply_daily_spend_identity(migrations_dir) return True except subprocess.TimeoutExpired: logger.warning( diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md index 05ffdf4f12f..164952d3edf 100644 --- a/litellm-proxy-extras/migration_runbook.md +++ b/litellm-proxy-extras/migration_runbook.md @@ -98,7 +98,7 @@ When the guard fires: `20260922000000_daily_spend_model_group_identity` separates daily usage by the requested public model group. Schedule a maintenance window: the six replacement unique indexes scan existing daily spend tables and block writes while building -Drain incoming requests and all in-memory and shared Redis spend buffers, then stop every old proxy and spend collector before applying this migration. Start only upgraded writers after it completes. Old writers require the removed conflict target and cannot run alongside the expanded identity. The migration changes indexes only and previously merged attribution stays as recorded. The new indexes and the writer's conflict target match `model_group` through `COALESCE("model_group", '')`, so rows stored with a NULL group before this change share one identity with the `''` new writers store and no row is rewritten. Prisma cannot declare an expression index, so `--use_prisma_db_push` deployments run this migration right after the push +Drain incoming requests and all in-memory and shared Redis spend buffers, then stop every old proxy and spend collector before applying this migration. Start only upgraded writers after it completes. Old writers require the removed conflict target and cannot run alongside the expanded identity. The migration changes indexes only; previously merged attribution and legacy NULL groups remain as recorded PTU flat charges retain their deployment/day identity across public-model renames. New inference rows include the requested group in their identity. Rolling back requires stopping writers and reconciling group-separated rows before restoring the old unique indexes; simply starting an older image is unsupported diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index 6d98f4037a4..c99410e3dd7 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -120,17 +120,6 @@ def _quoted(columns: Sequence[str]) -> str: return ", ".join(f'"{column}"' for column in columns) -_CONFLICT_EXPRESSIONS: Final[Mapping[str, str]] = MappingProxyType({"model_group": "COALESCE(\"model_group\", '')"}) - - -def _conflict_target(table: DailySpendTable) -> str: - """The unique index's expression list. model_group joined the key after rows were stored with - NULL in it, so the index matches it through COALESCE and those rows keep their identity unrewritten.""" - return ", ".join( - _CONFLICT_EXPRESSIONS.get(column, f'"{column}"') for column in (table.entity_id_column, *_KEY_COLUMNS) - ) - - def _as_text(value: object) -> str: return "" if value is None else str(value) @@ -230,7 +219,7 @@ def build_bulk_upsert( sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_conflict_target(table)}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f1a380356ce..cbd31d98814 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -873,6 +873,7 @@ model LiteLLM_DailyUserSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key") @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -941,6 +942,7 @@ model LiteLLM_DailyOrganizationSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key") @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -977,6 +979,7 @@ model LiteLLM_DailyEndUserSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key") @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -1013,6 +1016,7 @@ model LiteLLM_DailyAgentSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key") @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -1051,6 +1055,7 @@ model LiteLLM_DailyTeamSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key") @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -1089,6 +1094,7 @@ 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, model_group], map: "LiteLLM_DailyTagSpend_group_key") @@index([date]) @@index([tag, date]) @@index([api_key]) diff --git a/schema.prisma b/schema.prisma index f1a380356ce..cbd31d98814 100644 --- a/schema.prisma +++ b/schema.prisma @@ -873,6 +873,7 @@ model LiteLLM_DailyUserSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key") @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -941,6 +942,7 @@ model LiteLLM_DailyOrganizationSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key") @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -977,6 +979,7 @@ model LiteLLM_DailyEndUserSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key") @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -1013,6 +1016,7 @@ model LiteLLM_DailyAgentSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key") @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -1051,6 +1055,7 @@ model LiteLLM_DailyTeamSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt + @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key") @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -1089,6 +1094,7 @@ 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, model_group], map: "LiteLLM_DailyTagSpend_group_key") @@index([date]) @@index([tag, date]) @@index([api_key]) 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 be63305d6e3..ea4a25283a1 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 @@ -603,45 +603,6 @@ class TestResolveAllMigrationsLedger: assert self._executed_sql(calls) == _PARTITIONED_DRIFT_SQL -class TestDbPushAppliesDailySpendIdentity: - def _push(self, monkeypatch, tmp_path, run): - import litellm_proxy_extras.utils as utils_module - - calls = [] - - def fake_run(cmd, **kwargs): - calls.append(cmd) - return _FakeCompleted() - - monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) - monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fake_run) - monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {}) - monkeypatch.setattr(ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False)) - monkeypatch.setattr(ProxyExtrasDBManager, "raise_if_lens_rename_pending", staticmethod(lambda: None)) - assert run() is True - return calls - - @pytest.mark.parametrize( - "run", - [ - lambda: ProxyExtrasDBManager._run_migrations(use_migrate=False, use_v2_resolver=False), - lambda: ProxyExtrasDBManager._setup_database_v2(use_migrate=False), - ], - ids=["v1", "v2"], - ) - def test_the_identity_migration_runs_right_after_the_push(self, monkeypatch, tmp_path, run): - import litellm_proxy_extras.utils as utils_module - - calls = self._push(monkeypatch, tmp_path, run) - - assert [cmd[1:3] for cmd in calls] == [["db", "push"], ["db", "execute"]] - executed = calls[1] - assert executed[executed.index("--file") + 1] == ( - f"{tmp_path}/migrations/{utils_module.DAILY_SPEND_IDENTITY_MIGRATION}/migration.sql" - ) - assert executed[executed.index("--schema") + 1] == f"{tmp_path}/schema.prisma" - - class TestPartitionedSpendLogsPushGuard: def _forbid_subprocess(self, monkeypatch): import litellm_proxy_extras.utils as utils_module 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 de770e1b599..e470f93b16e 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -99,14 +99,13 @@ def test_one_statement_carries_every_row_in_the_batch(): assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1 -def test_conflict_target_is_the_unique_index_with_model_group_matched_through_coalesce(): +def test_conflict_target_is_the_full_unique_constraint(): sql, _ = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(),))) - conflict_target = re.search(r"ON CONFLICT \((.*)\) DO UPDATE", sql) + 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", ' - "COALESCE(\"model_group\", '')" + '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"' )