From d3c839147edf831efc6ccc83ce8d927f751d4077 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 15:21:35 +0000 Subject: [PATCH] fix(spend): keep CloudZero export and spend-log snapshots compatible with email recovery Co-authored-by: Mateo Wang --- litellm/integrations/cloudzero/database.py | 7 ++++++- litellm/integrations/focus/database.py | 7 ++++++- .../test_litellm/integrations/cloudzero/test_cloudzero.py | 5 +++++ .../spend_tracking/test_spend_management_endpoints.py | 1 + 4 files changed, 18 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 87f0c8bd160..6a63868a08b 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -98,9 +98,14 @@ class LiteLLMDatabase: fill_missing_api_key_aliases, ) + usage_rows: Final = ( + db_response.to_dicts() + if isinstance(db_response, pl.DataFrame) + else db_response if isinstance(db_response, list) else [] + ) # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. - recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) + recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as e: raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 96a32046e81..da6458c9369 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -100,9 +100,14 @@ class FocusLiteLLMDatabase: fill_missing_api_key_aliases, ) + usage_rows: Final = ( + db_response.to_dicts() + if isinstance(db_response, pl.DataFrame) + else db_response if isinstance(db_response, list) else [] + ) # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. - recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) + recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as exc: raise RuntimeError(f"Error retrieving usage data: {exc}") from exc diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py index 2d51eeb9944..c543156eedd 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -74,6 +74,8 @@ class TestCloudZeroHourlyExport: fake_db = MagicMock() async def query_raw_mock(query: str, *params): + if "LiteLLM_SpendLogs" in query: + return [] start_time_utc = params[0] if len(params) > 0 else None end_time_utc = params[1] if len(params) > 1 else None limit = params[2] if len(params) > 2 else None @@ -146,6 +148,9 @@ class TestCloudZeroHourlyExport: return joined fake_db.query_raw = AsyncMock(side_effect=query_raw_mock) + fake_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + fake_db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + fake_db.litellm_usertable.find_many = AsyncMock(return_value=[]) fake_client.db = fake_db mock_prisma_client_getter.return_value = fake_client diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 30b086bab61..b62c1c076c8 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -483,6 +483,7 @@ ignored_keys = [ "metadata.user_api_key_project_alias", "metadata.user_api_key_org_id", "metadata.user_api_key_user_id", + "metadata.user_api_key_user_email", "metadata.user_api_key_team_alias", "metadata.spend_logs_metadata", "metadata.requester_ip_address",