diff --git a/litellm/proxy/analytics_endpoints/analytics_endpoints.py b/litellm/proxy/analytics_endpoints/analytics_endpoints.py index f929cb74e40..4752593742c 100644 --- a/litellm/proxy/analytics_endpoints/analytics_endpoints.py +++ b/litellm/proxy/analytics_endpoints/analytics_endpoints.py @@ -84,7 +84,7 @@ async def get_global_activity( FROM "LiteLLM_SpendLogs" sl LEFT JOIN "LiteLLM_VerificationToken" vt ON sl."api_key" = vt."token" WHERE - sl."startTime" BETWEEN $1::date AND $2::date + interval '1 day' + sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') 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 2d3fc023a39..5be9d9bab3c 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -213,7 +213,7 @@ async def get_global_activity_internal_user( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND "user" = $3 GROUP BY date_trunc('day', "startTime") """ @@ -297,7 +297,7 @@ async def get_global_activity( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') GROUP BY date_trunc('day', "startTime") """ db_response = await prisma_client.db.query_raw( @@ -356,7 +356,7 @@ async def get_global_activity_model_internal_user( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND "user" = $3 GROUP BY model_group, date_trunc('day', "startTime") """ @@ -464,7 +464,7 @@ async def get_global_activity_model( COUNT(*) AS api_requests, SUM(total_tokens) AS total_tokens FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') GROUP BY model_group, date_trunc('day', "startTime") """ db_response = await prisma_client.db.query_raw( @@ -609,8 +609,7 @@ async def get_global_activity_exceptions_per_deployment( FROM "LiteLLM_ErrorLogs" WHERE - "startTime" >= $1::date - AND "startTime" < ($2::date + INTERVAL '1 day') + "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND model_group = $3 AND status_code = '429' GROUP BY @@ -741,8 +740,7 @@ async def get_global_activity_exceptions( FROM "LiteLLM_ErrorLogs" WHERE - "startTime" >= $1::date - AND "startTime" < ($2::date + INTERVAL '1 day') + "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND model_group = $3 AND status_code = '429' GROUP BY @@ -855,7 +853,7 @@ async def get_global_spend_provider( model_id, SUM(spend) AS spend FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date + WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND length(model_id) > 0 AND "user" = $3 GROUP BY model_id @@ -869,7 +867,7 @@ async def get_global_spend_provider( model_id, SUM(spend) AS spend FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date AND length(model_id) > 0 + WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND length(model_id) > 0 GROUP BY model_id """ db_response = await prisma_client.db.query_raw( @@ -1019,7 +1017,7 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" BETWEEN $1::date AND $2::date AND sl.api_key = $3 + sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND sl.api_key = $3 GROUP BY sl.api_key, sl.model @@ -1064,7 +1062,7 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" BETWEEN $1::date AND $2::date AND sl.user = $3 + sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND sl.user = $3 GROUP BY sl.api_key, sl.model @@ -1118,7 +1116,7 @@ async def get_global_spend_report( ON sl.team_id = tt.team_id WHERE - sl."startTime" BETWEEN $1::date AND $2::date + sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') GROUP BY date_trunc('day', sl."startTime"), tt.team_alias, @@ -1177,7 +1175,7 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" BETWEEN $1::date AND $2::date + sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') GROUP BY date_trunc('day', sl."startTime"), customer, @@ -1234,7 +1232,7 @@ async def get_global_spend_report( FROM "LiteLLM_SpendLogs" sl WHERE - sl."startTime" BETWEEN $1::date AND $2::date + sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') GROUP BY sl.api_key, sl.model @@ -1442,7 +1440,7 @@ async def _get_spend_report_for_time_range( jsonb_array_elements_text(request_tags) AS individual_request_tag, SUM(spend) AS total_spend FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::date AND "startTime" < ($2::date + INTERVAL '1 day') + WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') GROUP BY individual_request_tag ORDER BY total_spend DESC; """ @@ -2683,8 +2681,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::timestamp - AND "startTime" < $2::timestamp +WHERE "startTime" >= $1::timestamptz + AND "startTime" < $2::timestamptz 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 a5ad105c3ba..ec379df7114 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -483,7 +483,7 @@ async def get_spend_by_team_and_customer( ON sl.team_id = tt.team_id WHERE - sl."startTime" BETWEEN $1::date AND $2::date + sl."startTime" >= $1::timestamptz AND sl."startTime" < ($2::timestamptz + INTERVAL '1 day') 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 new file mode 100644 index 00000000000..f65c958b3db --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -0,0 +1,84 @@ +""" +Test that spend queries use timestamp filtering instead of date casting. + +This prevents the performance issue where date casting prevents index usage. +GitHub Issue: #17487 +""" + +import datetime +import os +import sys +from datetime import timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.spend_tracking.spend_tracking_utils import ( + get_spend_by_team_and_customer, +) + + +@pytest.mark.asyncio +async def test_spend_query_uses_timestamp_filtering(): + """ + Test that spend queries use timestamp filtering for index optimization. + + Verifies: + 1. SQL does NOT cast the startTime column to DATE (which prevents index usage) + 2. SQL uses >= and < operators with INTERVAL for timestamp range filtering + 3. Parameters passed are datetime objects (not date objects) + """ + # Mock prisma client + mock_prisma = MagicMock() + mock_db = MagicMock() + mock_query_raw = AsyncMock(return_value=[]) + mock_db.query_raw = mock_query_raw + mock_prisma.db = mock_db + + # Use timezone-aware datetime objects + start_date = datetime.datetime(2024, 1, 1, tzinfo=timezone.utc) + end_date = datetime.datetime(2024, 1, 31, tzinfo=timezone.utc) + + # Call the function + await get_spend_by_team_and_customer( + start_date=start_date, + end_date=end_date, + team_id="test_team", + customer_id="test_customer", + prisma_client=mock_prisma, + ) + + # Verify the query was called + assert mock_query_raw.called, "query_raw should have been called" + + # Extract SQL and parameters + # Prisma query_raw is called like: query_raw(sql, param1, param2, ...) + call_args = mock_query_raw.call_args[0] + sql = call_args[0] + 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" + + # 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" + + # 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"