fix(proxy-extras): build request-log indexes without psycopg, over socket hosts, and past a URL lock_timeout

Without psycopg the build warns and reports the indexes missing instead of crashing the job. Prisma's ?host= socket param now reaches libpq, so socket URLs connect, and the unreachable-host message names that socket. The build session sets lock_timeout to 0 next to statement_timeout, so a lock_timeout in the URL no longer cancels an online build that waits on an old snapshot

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -67,7 +67,11 @@ def _coordination_failure(database_url: str, exc: "psycopg.Error", *, connecting
         port: Final = parsed_url.port
     except ValueError:
         return prefix + f"P1001: Can't reach database server. {detail}"
-    host: Final = parsed_url.hostname or "localhost"
+    host: Final = (
+        next((value for key, value in parse_qsl(parsed_url.query) if key == "host"), None)
+        or parsed_url.hostname
+        or "localhost"
+    )
     return prefix + f"P1001: Can't reach database server at `{host}:{port or 5432}`. {detail}"
 
 

litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py

@@ -165,11 +165,19 @@ def ensure_request_log_indexes(
     replica gets in between partitions rather than waiting for the whole table. Any
     failure is logged and left for the next index build; the result says whether
     every index ended up valid. Never raises."""
-    import psycopg
+    try:
+        import psycopg
+    except ImportError:
+        logger.warning(
+            "psycopg is not installed; skipping the request-log index build. "
+            "Install the litellm[extra_proxy] extra, which includes psycopg."
+        )
+        return False
 
     try:
         with connect(database_url) as connection:
             connection.execute("SET statement_timeout = 0")
+            connection.execute("SET lock_timeout = 0")
             results: Final = tuple(_ensure_index(connection, schema, index) for index in indexes)
     except psycopg.Error as exc:
         logger.warning("Could not build the request-log indexes, leaving them for the next index build: %s", exc)

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -97,6 +97,7 @@ LIBPQ_URL_PARAMS: Final = frozenset(
         "gssencmode",
         "krbsrvname",
         "target_session_attrs",
+        "host",
     }
 )
 

tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

@@ -113,6 +113,12 @@ def test_coordination_failure_does_not_label_tls_failures_as_unreachable(
             "Name or service not known",
             "certificate-db.invalid",
         ),
+        (
+            "postgresql://u:pw@localhost:5432/db?host=/cloudsql/proj:us-central1:inst",
+            'connection failed: connection to server on socket "/cloudsql/proj:us-central1:inst/.s.PGSQL.5432" '
+            "failed: No such file or directory",
+            "/cloudsql/proj:us-central1:inst",
+        ),
     ],
 )
 def test_coordination_failure_labels_unreachable_hosts_named_like_tls_markers(
@@ -1548,6 +1554,15 @@ class TestStripPrismaQueryParams:
             == "postgresql://u:p@writer/db"
         )
 
+    def test_a_prisma_socket_host_param_reaches_libpq(self) -> None:
+        from psycopg.conninfo import conninfo_to_dict
+
+        url: Final = "postgresql://u:p@localhost/litellm?host=/cloudsql/proj:us-central1:inst&schema=tenant"
+
+        conninfo: Final = conninfo_to_dict(ProxyExtrasDBManager._strip_prisma_query_params(url))
+
+        assert conninfo["host"] == "/cloudsql/proj:us-central1:inst"
+
 
 class TestBuildRequestLogIndexes:
     """The migration job hands the index build the direct database URL and the schema
@@ -1592,6 +1607,17 @@ class TestBuildRequestLogIndexes:
 
         assert ProxyExtrasDBManager.build_request_log_indexes(build=lambda url, schema: False) is False
 
+    def test_without_psycopg_the_build_warns_and_reports_the_indexes_missing(
+        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        monkeypatch.setitem(sys.modules, "psycopg", None)
+        monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
+
+        with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
+            assert ProxyExtrasDBManager.build_request_log_indexes() is False
+
+        assert any("psycopg is not installed" in record.message for record in caplog.records)
+
     def test_without_a_database_url_nothing_is_built(
         self, monkeypatch: pytest.MonkeyPatch, builds: SimpleQueue[tuple[str, str]], build: Callable[[str, str], bool]
     ) -> None:

tests/proxy_migration_tests/test_request_log_indexes.py

@@ -26,6 +26,7 @@ from litellm_proxy_extras.request_log_indexes import (
 from litellm_proxy_extras.utils import ProxyExtrasDBManager
 from psycopg import sql
 from psycopg.abc import Params, QueryNoTemplate
+from psycopg.conninfo import make_conninfo
 from psycopg.rows import class_row
 
 pytestmark = pytest.mark.timeout(900)
@@ -436,6 +437,31 @@ def test_an_invalid_index_of_the_managed_name_on_a_plain_table_is_rebuilt(scratc
     assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
 
 
+def _close_a_second_after_a_build_waits_on_it(database_url: str, pin: "psycopg.Connection[tuple[object, ...]]") -> None:
+    deadline: Final = time.monotonic() + 60
+    with psycopg.connect(database_url, autocommit=True) as watcher:
+        while time.monotonic() < deadline and not watcher.execute(
+            "SELECT 1 FROM pg_stat_activity WHERE datname = current_database() AND wait_event = 'virtualxid'"
+        ).fetchone():
+            time.sleep(0.01)
+    time.sleep(1)
+    pin.close()
+
+
+@requires_db
+def test_a_lock_timeout_the_database_url_sets_does_not_cancel_the_online_build(scratch_database: str) -> None:
+    with_lock_timeout: Final = make_conninfo(scratch_database, options="-c lock_timeout=200")
+    pin: Final = _pin_a_snapshot_on(scratch_database, "LiteLLM_SpendLogs")
+    release_the_pin: Final = threading.Thread(target=_close_a_second_after_a_build_waits_on_it, args=(scratch_database, pin))
+    release_the_pin.start()
+    try:
+        assert ensure_request_log_indexes(with_lock_timeout, "public") is True
+    finally:
+        release_the_pin.join()
+
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+
+
 def _rebuild_as_another_replica(database_url: str, name: str, table: str, column: str) -> int:
     """Drop and rebuild the index from a second connection, as a replica that won the
     race would, and return the oid of the index it built."""