fix(spend): keep CloudZero export and spend-log snapshots compatible with email recovery

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-09-03 15:21:35 +00:00
parent 7e4032cfcc
commit d3c839147e
No known key found for this signature in database
4 changed files with 18 additions and 2 deletions

View file

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

View file

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

View file

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

View file

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