diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py index 2ae36bff1c0..f3d0f839251 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py +++ b/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}" 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 ec0e5368ab8..d75ce067906 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/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) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 75e2b00f538..399a15823b1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -97,6 +97,7 @@ LIBPQ_URL_PARAMS: Final = frozenset( "gssencmode", "krbsrvname", "target_session_attrs", + "host", } ) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 8907face226..70de3b9c0c7 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/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: diff --git a/tests/proxy_migration_tests/test_request_log_indexes.py b/tests/proxy_migration_tests/test_request_log_indexes.py index f8f34a39768..1b7fddd23c3 100644 --- a/tests/proxy_migration_tests/test_request_log_indexes.py +++ b/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."""