mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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.
This commit is contained in:
Mrinal Chanshetty
parent
a1971e4e52
commit
54ec7f7e51
2 changed files with 141 additions and 0 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue