diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 8b19576d86e..c51fc45f89e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -67,6 +67,8 @@ _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") _MIGRATION_DEADLOCK_MARKER = "deadlock detected" _MIGRATION_LOCK_TIMEOUT_MARKER = "canceling statement due to lock timeout" INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") +MCP_SERVER_TABLE: Final = "LiteLLM_MCPServerTable" +PROVIDER_URL_INDEX: Final = "LiteLLM_MCPServerTable_provider_url_key" _TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") _INVALID_LITELLM_INDEXES_SQL: Final = ( "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " @@ -940,6 +942,66 @@ class ProxyExtrasDBManager: return False return not remaining + @staticmethod + def ensure_provider_url_index(lock_timeout: str = "30s") -> bool: + """Create the partial unique index on provider MCP URLs when it is missing. + + The index exists only in migration SQL, because Prisma's schema cannot express a + partial index. `prisma db push` never creates it, and adopting a pushed database + marks migrations applied without running their SQL, so every database setup path + ends here. Never raises: returns True when the index exists, False when creation was + skipped or failed and will be retried on the next database setup run.""" + prisma_url: Final = os.getenv("DATABASE_URL") + if not prisma_url: + return False + try: + import psycopg + from psycopg import sql + except ImportError: + logger.warning( + "psycopg is not installed; skipping the provider MCP URL index check. " + "Install the litellm[extra_proxy] extra, which includes psycopg." + ) + return False + + schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) + index_name: Final = sql.Identifier(schema, PROVIDER_URL_INDEX) + try: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + present: Final = conn.execute("SELECT to_regclass(%s)", (index_name.as_string(conn),)).fetchone() + if present is not None and present[0] is not None: + return True + column: Final = conn.execute( + "SELECT 1 FROM information_schema.columns " + "WHERE table_schema = %s AND table_name = %s AND column_name = 'execution_mode'", + (schema, MCP_SERVER_TABLE), + ).fetchone() + if column is None: + return False + conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) + with held_migration_lock(conn) as held: + if not held: + logger.info( + "Another process holds the migration lock; leaving the provider MCP URL index " + "to the next database setup run" + ) + return False + conn.execute( + sql.SQL( + "CREATE UNIQUE INDEX IF NOT EXISTS {} ON {} (\"url\") WHERE \"execution_mode\" = 'provider'" + ).format(sql.Identifier(PROVIDER_URL_INDEX), sql.Identifier(schema, MCP_SERVER_TABLE)) + ) + except psycopg.Error as e: + logger.warning( + "Could not create the provider MCP URL unique index, will retry on the next database setup run. " + "Error: %s", + e, + ) + return False + logger.info("Created the provider MCP URL unique index that prisma db push and schema adoption skip") + return True + @staticmethod def _repair_under_migration_lock( conn: "psycopg.Connection[tuple[str, str, str]]", @@ -1422,6 +1484,7 @@ class ProxyExtrasDBManager: if not migrated: return False ProxyExtrasDBManager.repair_invalid_indexes() + ProxyExtrasDBManager.ensure_provider_url_index() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return True diff --git a/tests/integration/database/migrations/test_provider_url_index.py b/tests/integration/database/migrations/test_provider_url_index.py new file mode 100644 index 00000000000..94611463135 --- /dev/null +++ b/tests/integration/database/migrations/test_provider_url_index.py @@ -0,0 +1,78 @@ +import os +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from litellm_proxy_extras.utils import MCP_SERVER_TABLE, PROVIDER_URL_INDEX, ProxyExtrasDBManager + +psycopg = pytest.importorskip("psycopg") + +pytestmark = pytest.mark.timeout(120) + +requires_db: Final = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + + +def _base_url() -> str: + return os.environ["DATABASE_URL"].split("?")[0] + + +def _scratch_schema(monkeypatch: pytest.MonkeyPatch, columns: str) -> Iterator[str]: + schema: Final = f"provider_url_{uuid.uuid4().hex[:8]}" + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute(f'CREATE TABLE "{schema}"."{MCP_SERVER_TABLE}" ({columns})') + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") + yield schema + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + + +@pytest.fixture +def pushed_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + """What `prisma db push` or schema adoption leaves: the column, but not the SQL-only index.""" + yield from _scratch_schema(monkeypatch, 'server_id TEXT PRIMARY KEY, url TEXT, execution_mode TEXT NOT NULL') + + +@pytest.fixture +def older_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + yield from _scratch_schema(monkeypatch, "server_id TEXT PRIMARY KEY, url TEXT") + + +def _insert(schema: str, server_id: str, url: str, mode: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute( + f'INSERT INTO "{schema}"."{MCP_SERVER_TABLE}" (server_id, url, execution_mode) VALUES (%s, %s, %s)', + (server_id, url, mode), + ) + + +def _index_exists(schema: str) -> bool: + with psycopg.connect(_base_url(), autocommit=True) as conn: + row = conn.execute("SELECT to_regclass(%s)", (f'"{schema}"."{PROVIDER_URL_INDEX}"',)).fetchone() + return row is not None and row[0] is not None + + +@requires_db +def test_setup_creates_the_provider_url_index_a_pushed_database_lacks(pushed_schema: str) -> None: + assert _index_exists(pushed_schema) is False + + assert ProxyExtrasDBManager.ensure_provider_url_index() is True + assert ProxyExtrasDBManager.ensure_provider_url_index() is True + + assert _index_exists(pushed_schema) is True + _insert(pushed_schema, "gateway-a", "https://same.example/mcp", "litellm") + _insert(pushed_schema, "gateway-b", "https://same.example/mcp", "litellm") + _insert(pushed_schema, "provider-a", "https://same.example/mcp", "provider") + with pytest.raises(psycopg.errors.UniqueViolation): + _insert(pushed_schema, "provider-b", "https://same.example/mcp", "provider") + + +@requires_db +def test_setup_skips_the_provider_url_index_before_the_execution_mode_column_exists(older_schema: str) -> None: + assert ProxyExtrasDBManager.ensure_provider_url_index() is False + assert _index_exists(older_schema) is False