From 4d95818e34d22085af1be62ad6c4dad22c19aef7 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 2 Oct 2026 01:54:49 +0000 Subject: [PATCH] feat(proxy-extras): make the SpendLogs index build opt-in Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../request_log_indexes.py | 7 ++++--- .../litellm_proxy_extras/utils.py | 19 +++++++++++++---- .../test_request_log_indexes.py | 17 +++++++++++++++ .../test_litellm_proxy_extras_utils.py | 21 +++++++++++++++++++ 4 files changed, 57 insertions(+), 7 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 6c31e8364a9..3378368ffdf 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -1,6 +1,6 @@ -"""The request-log indexes built after `prisma migrate deploy` instead of by a migration: -by the migration job, or by a serving proxy that ran the migrations itself (in the -background, once it serves). +"""The request-log indexes are built after `prisma migrate deploy` instead of by a migration, +but only when `LITELLM_BUILD_SPEND_LOGS_INDEXES` is set: by the migration job, or by a serving +proxy that ran the migrations itself (in the background, once it serves). A migration cannot build them: a plain `CREATE INDEX` blocks spend-log inserts for the whole build, and `CREATE INDEX CONCURRENTLY` is refused on a partitioned parent @@ -50,6 +50,7 @@ class RequestLogIndex: return f"{kept}_{digest}" +REQUEST_LOG_INDEXES_ENV_VAR: Final = "LITELLM_BUILD_SPEND_LOGS_INDEXES" REQUEST_LOG_INDEXES: Final = ( RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'), RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'), diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 3acc19d397d..fbe2c784c75 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -26,7 +26,11 @@ from litellm_proxy_extras.replica_identity import ( REPLICA_IDENTITY_FULL_ENV_VAR, apply_replica_identity_full, ) -from litellm_proxy_extras.request_log_indexes import ensure_request_log_indexes, filter_request_log_index_diff +from litellm_proxy_extras.request_log_indexes import ( + REQUEST_LOG_INDEXES_ENV_VAR, + ensure_request_log_indexes, + filter_request_log_index_diff, +) if TYPE_CHECKING: import psycopg @@ -1263,9 +1267,16 @@ class ProxyExtrasDBManager: @staticmethod def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool: """Build the indexes in `REQUEST_LOG_INDEXES` on the writer, in the schema the - migrations target. Idempotent and never raises; False when an index is still - missing or invalid, so the migration job reports it and gets rerun instead of - leaving the table unindexed until the next deploy.""" + migrations target. Does nothing unless `LITELLM_BUILD_SPEND_LOGS_INDEXES` is set. + Idempotent and never raises; False when an index is still missing or invalid, so + the migration job reports it and gets rerun instead of leaving the table unindexed + until the next deploy.""" + if not str_to_bool(os.getenv(REQUEST_LOG_INDEXES_ENV_VAR)): + logger.info( + "%s is not enabled, skipping the request-log index build", + REQUEST_LOG_INDEXES_ENV_VAR, + ) + return True database_url: Final = os.environ.get("DATABASE_URL") if not database_url: return True diff --git a/tests/proxy_migration_tests/test_request_log_indexes.py b/tests/proxy_migration_tests/test_request_log_indexes.py index 23e4adce477..cfc74e1a501 100644 --- a/tests/proxy_migration_tests/test_request_log_indexes.py +++ b/tests/proxy_migration_tests/test_request_log_indexes.py @@ -19,6 +19,7 @@ from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, migration_lo from litellm_proxy_extras.migration_recovery import roll_back_failed_inert_migration from litellm_proxy_extras.request_log_indexes import ( REQUEST_LOG_INDEXES, + REQUEST_LOG_INDEXES_ENV_VAR, RequestLogIndex, build_index_on_partitioned_table, ensure_request_log_indexes, @@ -149,6 +150,7 @@ def scratch_database(release: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Pa _deploy_release(database_url, tmp_path / "prisma", release) monkeypatch.delenv("DIRECT_URL", raising=False) monkeypatch.setenv("DATABASE_URL", database_url) + monkeypatch.setenv(REQUEST_LOG_INDEXES_ENV_VAR, "true") yield database_url finally: with psycopg.connect(admin_url, autocommit=True) as conn: @@ -583,6 +585,21 @@ def test_a_migration_job_that_could_not_build_the_indexes_reports_failure_and_su assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True} +@requires_db +@RESOLVERS +def test_a_migration_job_without_the_opt_in_succeeds_and_builds_no_index( + scratch_database: str, use_v2_resolver: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + api_key_index_before: Final = _index_validity(scratch_database, "api_key_startTime_idx") + monkeypatch.delenv(REQUEST_LOG_INDEXES_ENV_VAR, raising=False) + + assert _migration_job(use_v2_resolver) is True + + _assert_every_ledger_row_is_finished(scratch_database) + assert _index_validity(scratch_database, "litellm_call_id_idx") == {} + assert _index_validity(scratch_database, "api_key_startTime_idx") == api_key_index_before + + @requires_db @RESOLVERS def test_the_serving_proxy_setup_applies_the_inert_migrations_and_builds_no_index( diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 5540cf54193..44d0c474525 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -15,6 +15,7 @@ sys.path.insert( ), ) +from litellm_proxy_extras.request_log_indexes import REQUEST_LOG_INDEXES_ENV_VAR from litellm_proxy_extras.utils import ( PARTITIONED_SPEND_LOGS_PUSH_ERROR, ProxyExtrasDBManager, @@ -1033,6 +1034,10 @@ class TestBuildRequestLogIndexes: """The migration job hands the index build the direct database URL and the schema the migrations target, waits for it, and reports its result.""" + @pytest.fixture(autouse=True) + def _enable_request_log_index_build(self, monkeypatch): + monkeypatch.setenv(REQUEST_LOG_INDEXES_ENV_VAR, "true") + @pytest.fixture def builds(self): return [] @@ -1073,6 +1078,22 @@ class TestBuildRequestLogIndexes: assert builds == [] + def test_without_the_opt_in_nothing_is_built(self, monkeypatch, builds, build): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db") + monkeypatch.delenv(REQUEST_LOG_INDEXES_ENV_VAR, raising=False) + + assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True + + assert builds == [] + + def test_an_opt_in_of_false_builds_nothing(self, monkeypatch, builds, build): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db") + monkeypatch.setenv(REQUEST_LOG_INDEXES_ENV_VAR, "false") + + assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True + + assert builds == [] + class TestStartRequestLogIndexBuild: """A serving proxy that ran the migrations starts the index build on a daemon thread