From 86ff6f994c3ec346bb6442b35387ce6beaaf7161 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 17:50:49 -0800 Subject: [PATCH] revert(mcp): drop startup creation of the provider URL index Duplicate provider URLs are now rejected on write by an app-level check under an advisory lock, and enforcement considers every registration of a URL, so nothing depends on the partial unique index. The index stays in the regular migration as an extra guard; startup no longer runs DDL outside prisma migrate. This reverts commit 54ec7f7e51. --- .../litellm_proxy_extras/utils.py | 63 --------------- .../migrations/test_provider_url_index.py | 78 ------------------- 2 files changed, 141 deletions(-) delete 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 c51fc45f89e..8b19576d86e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -67,8 +67,6 @@ _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)) " @@ -942,66 +940,6 @@ 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]]", @@ -1484,7 +1422,6 @@ 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 deleted file mode 100644 index 94611463135..00000000000 --- a/tests/integration/database/migrations/test_provider_url_index.py +++ /dev/null @@ -1,78 +0,0 @@ -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