Merge 4d95818e34 into d80f8c28ca

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")'),

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -28,7 +28,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
@@ -1380,9 +1384,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

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(

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -18,6 +18,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,
@@ -1283,6 +1284,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 []
@@ -1323,6 +1328,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