scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

fix(sql): Optimize SpendLogs queries to use timestamp filtering for index usage (#17504)

Browse source 
* fix: optimize SpendLogs queries to use timestamp filtering (#17487)

* use timestamptz & enhance test
This commit is contained in:
Kristian Brünn 2025-12-05 06:52:57 +01:00 • committed by GitHub
parent 8776336c3c
commit 63fae79493
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 102 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/analytics_endpoints/analytics_endpoints.py
View file

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

34
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

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

2
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

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

84
tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py Normal file
View file

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