mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
2310daef7e
Through a transaction-mode pooler Prisma takes pg_advisory_lock(72707369) on one backend and releases it on another, so the first backend keeps the lock while idle and the next boot waits on it. Migrations now run with PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK=1 while LiteLLM holds the coordinator lock and Prisma's key as transaction-level locks on one pinned connection, so nothing can outlive the migration Resolves LIT-9571 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
58 lines
2.4 KiB
Python
58 lines
2.4 KiB
Python
import os
|
|
import secrets
|
|
from typing import Final
|
|
|
|
import psycopg
|
|
import pytest
|
|
from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, PRISMA_ADVISORY_LOCK_KEY, prisma_migration_lock
|
|
|
|
pytestmark = pytest.mark.skipif("DATABASE_URL" not in os.environ, reason="needs a live Postgres")
|
|
|
|
|
|
def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
|
|
row: Final = connection.execute("SELECT pg_try_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)).fetchone()
|
|
taken: Final = row is not None and row[0] is True
|
|
if taken:
|
|
connection.execute("SELECT pg_advisory_unlock(%s)", (PRISMA_ADVISORY_LOCK_KEY,))
|
|
return taken
|
|
|
|
|
|
def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]:
|
|
rows: Final = connection.execute(
|
|
"SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted"
|
|
).fetchall()
|
|
return frozenset(int(row[0]) for row in rows)
|
|
|
|
|
|
def _index_builds_concurrently_now(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
|
|
table: Final = f'"lit_cic_{secrets.token_hex(4)}"'
|
|
connection.execute(f"CREATE TABLE {table} (a int)")
|
|
try:
|
|
connection.execute("SET lock_timeout = '5s'")
|
|
connection.execute(f"CREATE INDEX CONCURRENTLY ON {table} (a)")
|
|
except psycopg.errors.LockNotAvailable:
|
|
return False
|
|
finally:
|
|
connection.execute("RESET lock_timeout")
|
|
connection.execute(f"DROP TABLE {table}")
|
|
return True
|
|
|
|
|
|
def test_prisma_migration_lock_does_not_stall_a_concurrent_index_build():
|
|
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
|
|
with prisma_migration_lock(os.environ["DATABASE_URL"]):
|
|
assert _index_builds_concurrently_now(peer)
|
|
|
|
|
|
def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leaves_nothing_behind():
|
|
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
|
|
assert _can_take_prismas_session_lock(peer)
|
|
|
|
with prisma_migration_lock(os.environ["DATABASE_URL"]):
|
|
held: Final = _granted_advisory_keys(peer)
|
|
assert PRISMA_ADVISORY_LOCK_KEY in held
|
|
assert MIGRATION_LOCK_KEY in held
|
|
assert not _can_take_prismas_session_lock(peer)
|
|
|
|
assert _granted_advisory_keys(peer).isdisjoint({PRISMA_ADVISORY_LOCK_KEY, MIGRATION_LOCK_KEY})
|
|
assert _can_take_prismas_session_lock(peer)
|