From c7934c460d4182f5545c385429d42ae3da7faf0e Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 10 Apr 2026 17:04:52 -0700 Subject: [PATCH] fix(spend): session-TZ-independent date filtering for spend/error log queries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The raw SQL queries in spend_management_endpoints.py, spend_tracking_utils.py, and analytics_endpoints.py cast date params to ::timestamptz while comparing against the plain-timestamp "startTime" column. Postgres resolves the type mismatch by promoting the column using the DB session timezone, which drifts the filter window and date_trunc buckets whenever session TZ is not UTC — silently dropping rows at UTC day boundaries and, for narrow windows, losing rows entirely. Wraps every such comparison with `AT TIME ZONE 'UTC'` so the param side resolves to a plain timestamp matching the column type. Both sides end up as plain timestamp, Postgres does no implicit conversion, and session TZ plays no role in the query. The fix is constant-foldable so the existing startTime index (PR #17504) remains usable. Also marks strptime-produced datetimes as tz-aware UTC at the call sites for intent clarity and consistency with parse_date. Fixes #22529 (Logs page missing recent rows under non-UTC session TZ) and the Global Usage single-day-returns-two-days symptom reported internally. --- .../analytics_endpoints.py | 13 +- .../spend_management_endpoints.py | 99 ++++++---- .../spend_tracking/spend_tracking_utils.py | 3 +- .../test_spend_query_optimization.py | 176 ++++++++++++++++-- 4 files changed, 235 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/analytics_endpoints/analytics_endpoints.py b/litellm/proxy/analytics_endpoints/analytics_endpoints.py index 4752593742c..6835f0c9095 100644 --- a/litellm/proxy/analytics_endpoints/analytics_endpoints.py +++ b/litellm/proxy/analytics_endpoints/analytics_endpoints.py @@ -1,5 +1,5 @@ #### Analytics Endpoints ##### -from datetime import datetime +from datetime import datetime, timezone from typing import List, Optional import fastapi @@ -58,8 +58,10 @@ async def get_global_activity( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -83,8 +85,9 @@ async def get_global_activity( SUM(CASE WHEN sl."cache_hit" != 'True' THEN sl."completion_tokens" ELSE 0 END) AS generated_completion_tokens FROM "LiteLLM_SpendLogs" sl LEFT JOIN "LiteLLM_VerificationToken" vt ON sl."api_key" = vt."token" - WHERE - sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + WHERE + sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY vt."key_alias", sl."call_type", diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 3c1b7cfd10c..de51f3b0a84 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -223,7 +223,8 @@ async def get_global_activity_internal_user( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') AND "user" = $3 GROUP BY date_trunc('day', "startTime") """ @@ -282,8 +283,10 @@ async def get_global_activity( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -307,7 +310,8 @@ async def get_global_activity( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY date_trunc('day', "startTime") """ db_response = await prisma_client.db.query_raw( @@ -366,7 +370,8 @@ async def get_global_activity_model_internal_user( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') AND "user" = $3 GROUP BY model_group, date_trunc('day', "startTime") """ @@ -448,8 +453,10 @@ async def get_global_activity_model( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -474,7 +481,8 @@ async def get_global_activity_model( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY model_group, date_trunc('day', "startTime") """ db_response = await prisma_client.db.query_raw( @@ -600,8 +608,10 @@ async def get_global_activity_exceptions_per_deployment( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -619,7 +629,8 @@ async def get_global_activity_exceptions_per_deployment( FROM "LiteLLM_ErrorLogs" WHERE - "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') AND model_group = $3 AND status_code = '429' GROUP BY @@ -732,8 +743,10 @@ async def get_global_activity_exceptions( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -750,7 +763,8 @@ async def get_global_activity_exceptions( FROM "LiteLLM_ErrorLogs" WHERE - "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') AND model_group = $3 AND status_code = '429' GROUP BY @@ -837,8 +851,10 @@ async def get_global_spend_provider( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import llm_router, prisma_client @@ -863,7 +879,8 @@ async def get_global_spend_provider( model_id, SUM(spend) AS spend FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') AND length(model_id) > 0 AND "user" = $3 GROUP BY model_id @@ -877,7 +894,9 @@ async def get_global_spend_provider( model_id, SUM(spend) AS spend FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND length(model_id) > 0 + WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') + AND length(model_id) > 0 GROUP BY model_id """ db_response = await prisma_client.db.query_raw( @@ -996,8 +1015,10 @@ async def get_global_spend_report( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( + tzinfo=timezone.utc + ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import premium_user, prisma_client @@ -1029,7 +1050,9 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND sl.api_key = $3 + sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') + AND sl.api_key = $3 GROUP BY sl.api_key, sl.model @@ -1074,7 +1097,9 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND sl.user = $3 + sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') + AND sl.user = $3 GROUP BY sl.api_key, sl.model @@ -1128,7 +1153,8 @@ async def get_global_spend_report( ON sl.team_id = tt.team_id WHERE - sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY date_trunc('day', sl."startTime"), tt.team_alias, @@ -1187,7 +1213,8 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY date_trunc('day', sl."startTime"), customer, @@ -1244,7 +1271,8 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY sl.api_key, sl.model @@ -1448,11 +1476,12 @@ async def _get_spend_report_for_time_range( # get spend per tag for today sql_query = """ - SELECT + SELECT jsonb_array_elements_text(request_tags) AS individual_request_tag, SUM(spend) AS total_spend FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') + WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY individual_request_tag ORDER BY total_spend DESC; """ @@ -1910,11 +1939,17 @@ async def ui_view_spend_logs( # noqa: PLR0915 sql_params: List[Any] = [] p = 1 # parameter index counter - # Date range (always present) - sql_conditions.append(f'"startTime" >= ${p}::timestamptz') + # Date range (always present). Wrap the param side with + # `AT TIME ZONE 'UTC'` so comparison against the plain `timestamp` + # column does not depend on the DB session timezone (see #22529). + sql_conditions.append( + f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')" + ) sql_params.append(start_date_obj) p += 1 - sql_conditions.append(f'"startTime" <= ${p}::timestamptz') + sql_conditions.append( + f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')" + ) sql_params.append(end_date_obj) p += 1 @@ -2897,8 +2932,8 @@ async def global_spend_end_users(data: Optional[GlobalEndUsersSpend] = None): sql_query = """ SELECT end_user, COUNT(*) AS total_count, SUM(spend) AS total_spend FROM "LiteLLM_SpendLogs" -WHERE "startTime" >= $1::timestamptz - AND "startTime" < $2::timestamptz +WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND "startTime" < ($2::timestamptz AT TIME ZONE 'UTC') AND ( CASE WHEN $3::TEXT IS NULL THEN TRUE diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8ea93453ed3..8a963ec0134 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -559,7 +559,8 @@ async def get_spend_by_team_and_customer( ON sl.team_id = tt.team_id WHERE - sl."startTime" >= $1::timestamptz AND sl."startTime" < ($2::timestamptz + INTERVAL '1 day') + sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') AND sl.team_id = $3 AND sl.end_user = $4 GROUP BY diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index f65c958b3db..fbac71e6372 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -60,25 +60,165 @@ async def test_spend_query_uses_timestamp_filtering(): params = call_args[1:] # 1) SQL should NOT cast the startTime column to DATE (prevents index usage) - assert "::date" not in sql.lower(), \ - "SQL should not use '::date' casting which prevents index usage" - assert "date(" not in sql.lower(), \ - "SQL should not use DATE() function which prevents index usage" + assert ( + "::date" not in sql.lower() + ), "SQL should not use '::date' casting which prevents index usage" + assert ( + "date(" not in sql.lower() + ), "SQL should not use DATE() function which prevents index usage" # 2) SQL should use timestamp-range filtering pattern for index optimization - assert '"startTime" >=' in sql or '"startTime">=' in sql, \ - "SQL should use >= operator for lower bound" - assert '"startTime" <' in sql or '"startTime"<' in sql, \ - "SQL should use < operator for upper bound" - assert "interval '1 day'" in sql.lower(), \ - "SQL should use INTERVAL for date arithmetic" + assert ( + '"startTime" >=' in sql or '"startTime">=' in sql + ), "SQL should use >= operator for lower bound" + assert ( + '"startTime" <' in sql or '"startTime"<' in sql + ), "SQL should use < operator for upper bound" + assert ( + "interval '1 day'" in sql.lower() + ), "SQL should use INTERVAL for date arithmetic" # 3) Parameters should be datetime objects (not date objects) - assert isinstance(params[0], datetime.datetime), \ - "First parameter (start_date) should be datetime object" - assert isinstance(params[1], datetime.datetime), \ - "Second parameter (end_date) should be datetime object" - assert params[0].tzinfo is not None, \ - "start_date should be timezone-aware" - assert params[1].tzinfo is not None, \ - "end_date should be timezone-aware" + assert isinstance( + params[0], datetime.datetime + ), "First parameter (start_date) should be datetime object" + assert isinstance( + params[1], datetime.datetime + ), "Second parameter (end_date) should be datetime object" + assert params[0].tzinfo is not None, "start_date should be timezone-aware" + assert params[1].tzinfo is not None, "end_date should be timezone-aware" + + +@pytest.mark.asyncio +async def test_global_activity_wraps_params_in_at_time_zone_utc(monkeypatch): + """ + /global/activity must emit `AT TIME ZONE 'UTC'` around its date params + so the date window and `date_trunc` bucketing do not depend on the DB + session timezone. Regression guard for Issue 1. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + get_global_activity, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + + await get_global_activity( + start_date="2026-02-16", + end_date="2026-02-16", + user_api_key_dict=auth, + ) + + assert mock_prisma.db.query_raw.called, "query_raw should have been called" + call_args = mock_prisma.db.query_raw.call_args[0] + sql = call_args[0] + params = call_args[1:] + + # 1) SQL must wrap both bounds in `AT TIME ZONE 'UTC'`. + assert sql.count("AT TIME ZONE 'UTC'") >= 2, ( + "Both date bounds must be wrapped with `AT TIME ZONE 'UTC'` so that " + "comparison against the plain-timestamp column is session-TZ-independent. " + f"SQL was:\n{sql}" + ) + + # 2) Params must still be tz-aware UTC datetimes (preserves existing contract). + assert isinstance(params[0], datetime.datetime) + assert isinstance(params[1], datetime.datetime) + assert params[0].tzinfo is not None and params[0].utcoffset() == datetime.timedelta( + 0 + ) + assert params[1].tzinfo is not None and params[1].utcoffset() == datetime.timedelta( + 0 + ) + + +@pytest.mark.asyncio +async def test_global_activity_internal_user_wraps_params_in_at_time_zone_utc( + monkeypatch, +): + """ + The internal-user branch of /global/activity goes through a different + helper (`get_global_activity_internal_user`) and has its own SQL string. + Both branches must carry the fix. Regression guard for Issue 1. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + get_global_activity, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1" + ) + + await get_global_activity( + start_date="2026-02-16", + end_date="2026-02-16", + user_api_key_dict=auth, + ) + + assert mock_prisma.db.query_raw.called + sql = mock_prisma.db.query_raw.call_args[0][0] + assert sql.count("AT TIME ZONE 'UTC'") >= 2, ( + "Internal-user branch must also wrap date bounds with " + f"`AT TIME ZONE 'UTC'`. SQL was:\n{sql}" + ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): + """ + /spend/logs/ui builds its WHERE clause dynamically. The date-range + conditions must wrap the param side with `AT TIME ZONE 'UTC'` so the + log filter window doesn't drift with the DB session TZ. Regression + guard for GH #22529. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + assert mock_prisma.db.query_raw.called, "query_raw should have been called" + sql = mock_prisma.db.query_raw.call_args[0][0] + assert sql.count("AT TIME ZONE 'UTC'") >= 2, ( + "/spend/logs/ui must wrap both `startTime` bounds with " + f"`AT TIME ZONE 'UTC'`. SQL was:\n{sql}" + )