fix(migrations): exit the migration job by the migrations' result

An index the job could not build, for example while another replica held the migration lock, failed a helm upgrade whose migrations had passed. The job still waits for the build and logs what is left, and the next run or a serving proxy builds it

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -1303,8 +1303,7 @@ class ProxyExtrasDBManager:
     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."""
+        missing or invalid, which the builder logs for the next index build to pick up."""
         database_url: Final = os.environ.get("DATABASE_URL")
         if not database_url:
             return True
@@ -1322,9 +1321,13 @@ class ProxyExtrasDBManager:
         build: Callable[[], bool] = build_request_log_indexes,
     ) -> bool:
         """The migration job's whole run: `setup_database`, then the request-log indexes,
-        built synchronously so the job exits only once they are in place. False when the
-        migrations failed or an index could not be built, so the Job is rerun."""
-        return setup(use_migrate, use_v2_resolver) and build()
+        built synchronously so the job exits only once the build has finished. The result is
+        the migrations' alone, so an index that could not be built never fails an upgrade
+        whose migrations passed."""
+        if not setup(use_migrate, use_v2_resolver):
+            return False
+        build()
+        return True
 
     @staticmethod
     def start_request_log_index_build(build: Callable[[], bool] = build_request_log_indexes) -> threading.Thread:

migrations/run.py

@@ -5,9 +5,10 @@ recovery logic in `litellm_proxy_extras.ProxyExtrasDBManager.setup_database`
 (P3005 baseline + P3009/P3018 idempotent-error handling, retries, etc.), then
 builds the request-log indexes the migrations leave out
 (`litellm_proxy_extras.request_log_indexes`), waiting for them. The job exits
-non-zero when an index could not be built so that it is rerun. A serving proxy
-that runs the migrations itself builds the same indexes in the background once
-it serves.
+by the migrations' result, so an index that could not be built is logged and
+left to the next run instead of failing an upgrade whose migrations passed. A
+serving proxy that runs the migrations itself builds the same indexes in the
+background once it serves.
 
 Env vars:
   DATABASE_URL                  required unless it can be assembled at

tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

@@ -1587,9 +1587,7 @@ class TestBuildRequestLogIndexes:
         assert builds.get_nowait() == ("postgresql://u:p@primary:5432/db", "public")
         assert builds.empty()
 
-    def test_a_build_that_leaves_indexes_missing_is_reported_so_the_job_reruns(
-        self, monkeypatch: pytest.MonkeyPatch
-    ) -> None:
+    def test_a_build_that_leaves_indexes_missing_is_reported(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
 
         assert ProxyExtrasDBManager.build_request_log_indexes(build=lambda url, schema: False) is False
@@ -1628,8 +1626,8 @@ class TestStartRequestLogIndexBuild:
 
 
 class TestRunMigrationJob:
-    """`run_migration_job` is `setup_database` followed by the index build, each step's
-    result deciding whether the job reports success."""
+    """`run_migration_job` is `setup_database` followed by the index build, the migrations'
+    result alone deciding whether the job reports success."""
 
     @pytest.fixture
     def calls(self) -> SimpleQueue[tuple[str, bool, bool] | tuple[str]]:
@@ -1681,13 +1679,13 @@ class TestRunMigrationJob:
         assert calls.get_nowait() == ("setup", True, True)
         assert calls.empty()
 
-    def test_the_job_fails_when_an_index_could_not_be_built(
+    def test_the_job_succeeds_by_the_migrations_when_an_index_could_not_be_built(
         self,
         calls: SimpleQueue[tuple[str, bool, bool] | tuple[str]],
         setup: Callable[[bool], Callable[[bool, bool], bool]],
         build: Callable[[bool], Callable[[], bool]],
     ) -> None:
-        assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(True), build=build(False)) is False
+        assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(True), build=build(False)) is True
 
         assert (calls.get_nowait(), calls.get_nowait()) == (("setup", True, True), ("build",))
         assert calls.empty()

tests/proxy_migration_tests/test_request_log_indexes.py

@@ -581,15 +581,14 @@ def test_a_process_holding_the_migration_lock_makes_the_build_wait_for_the_next_
 
 @requires_db
 @RESOLVERS
-def test_a_migration_job_that_could_not_build_the_indexes_reports_failure_and_succeeds_when_rerun(
+def test_a_migration_job_that_could_not_build_the_indexes_succeeds_and_a_rerun_builds_them(
     scratch_database: str, use_v2_resolver: bool
 ) -> None:
-    """The migration job waits for the build and exits by run_migration_job's result; a job
-    that exits 0 with the indexes missing would leave the table unindexed until the next
-    deploy or until a serving proxy's background build gets to them."""
+    """The migration job exits by the migrations' result, so an index another replica's lock kept
+    it from building never fails an upgrade, and the next run builds it."""
     with psycopg.connect(scratch_database, autocommit=True) as other_replica:
         other_replica.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
-        assert _migration_job(use_v2_resolver) is 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") == {}