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 bc803c0a096..3f52ce49dc0 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", "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", 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_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_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_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_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_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_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_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_DailyTagSpend_group_key" ON "LiteLLM_DailyTagSpend" ("tag", "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", '')); 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 cbd31d98814..f1a380356ce 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -873,7 +873,6 @@ 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]) @@ -942,7 +941,6 @@ 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]) @@ -979,7 +977,6 @@ 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]) @@ -1016,7 +1013,6 @@ 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]) @@ -1055,7 +1051,6 @@ 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]) @@ -1094,7 +1089,6 @@ 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 245244250ee..abacaf7ca39 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -202,6 +202,9 @@ 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")): @@ -1027,6 +1030,7 @@ class ProxyExtrasDBManager: stdout=None, stderr=None, ) + ProxyExtrasDBManager._apply_daily_spend_identity(migrations_dir) return True except ( subprocess.CalledProcessError, @@ -1351,6 +1355,24 @@ 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: @@ -1584,6 +1606,7 @@ 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 164952d3edf..05ffdf4f12f 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; previously merged attribution and legacy NULL groups remain as recorded +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 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 c99410e3dd7..6d98f4037a4 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -120,6 +120,17 @@ 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) @@ -219,7 +230,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 ({_conflict_target(table)}) 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 cbd31d98814..f1a380356ce 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -873,7 +873,6 @@ 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]) @@ -942,7 +941,6 @@ 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]) @@ -979,7 +977,6 @@ 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]) @@ -1016,7 +1013,6 @@ 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]) @@ -1055,7 +1051,6 @@ 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]) @@ -1094,7 +1089,6 @@ 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 cbd31d98814..f1a380356ce 100644 --- a/schema.prisma +++ b/schema.prisma @@ -873,7 +873,6 @@ 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]) @@ -942,7 +941,6 @@ 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]) @@ -979,7 +977,6 @@ 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]) @@ -1016,7 +1013,6 @@ 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]) @@ -1055,7 +1051,6 @@ 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]) @@ -1094,7 +1089,6 @@ 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 ea4a25283a1..be63305d6e3 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,6 +603,45 @@ 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 e470f93b16e..de770e1b599 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -99,13 +99,14 @@ 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_full_unique_constraint(): +def test_conflict_target_is_the_unique_index_with_model_group_matched_through_coalesce(): sql, _ = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(),))) - conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) + conflict_target = re.search(r"ON CONFLICT \((.*)\) DO UPDATE", 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", "model_group"' + '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", ' + "COALESCE(\"model_group\", '')" )