From 780db6a4c8d9902af5e512d24563b5da8a5925ce Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 16:57:59 -0800 Subject: [PATCH] fix(mcp): create the provider URL unique index after every database setup The partial unique index on provider MCP URLs exists only in migration SQL. prisma db push never creates it, and adopting a pushed database marks migrations applied without running their SQL, so those setups accepted duplicate provider URLs. Database setup now creates the index when it is missing, under the migration lock, and never fails startup. --- .../litellm_proxy_extras/utils.py | 63 +++++++++++++++ .../migrations/test_provider_url_index.py | 78 +++++++++++++++++++ 2 files changed, 141 insertions(+) create mode 100644 tests/integration/database/migrations/test_provider_url_index.py 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