mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge pull request #39384 from BerriAI/litellm_fix_invalid_index_after_migration_deadlock
fix(proxy-extras): rebuild indexes left INVALID by a migration deadlock
This commit is contained in:
Mateo Wang
GitHub
commit
82e6b84f5a
2 changed files with 387 additions and 7 deletions
|
|
@ -8,14 +8,10 @@ import tempfile
|
|||
import time
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.replica_identity import (
|
||||
REPLICA_IDENTITY_FULL_ENV_VAR,
|
||||
apply_replica_identity_full,
|
||||
)
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
|
||||
|
|
@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
|
|||
prisma_command_timeout,
|
||||
prisma_migrate_deploy_timeout,
|
||||
)
|
||||
from litellm_proxy_extras.replica_identity import (
|
||||
REPLICA_IDENTITY_FULL_ENV_VAR,
|
||||
apply_replica_identity_full,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
import psycopg.sql
|
||||
|
||||
|
||||
def str_to_bool(value: Optional[str]) -> bool:
|
||||
|
|
@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
|
|||
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
|
||||
|
||||
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
|
||||
INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
|
||||
_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)) "
|
||||
"FROM pg_index i "
|
||||
"JOIN pg_class c ON c.oid = i.indexrelid "
|
||||
"JOIN pg_class t ON t.oid = i.indrelid "
|
||||
"JOIN pg_namespace n ON n.oid = t.relnamespace "
|
||||
"WHERE NOT i.indisvalid "
|
||||
" AND c.relkind = 'i' "
|
||||
" AND n.nspname = %s "
|
||||
" AND t.relname LIKE %s "
|
||||
" AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) "
|
||||
"ORDER BY c.relname"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _InvalidIndex:
|
||||
schema: str
|
||||
name: str
|
||||
table_size: str
|
||||
|
||||
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
|
||||
|
||||
|
|
@ -624,7 +650,7 @@ class ProxyExtrasDBManager:
|
|||
def _strip_prisma_query_params(url: str) -> str:
|
||||
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
|
||||
schema, etc.) from DATABASE_URL so psycopg can parse it."""
|
||||
from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
|
||||
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
|
||||
|
||||
parsed = urlparse(url)
|
||||
if not parsed.query:
|
||||
|
|
@ -645,7 +671,7 @@ class ProxyExtrasDBManager:
|
|||
"target_session_attrs",
|
||||
}
|
||||
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
|
||||
return urlunparse(parsed._replace(query=urlencode(kept)))
|
||||
return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
|
||||
|
||||
@staticmethod
|
||||
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
|
||||
|
|
@ -719,6 +745,95 @@ class ProxyExtrasDBManager:
|
|||
", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _invalid_litellm_indexes(
|
||||
conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
|
||||
) -> tuple[_InvalidIndex, ...]:
|
||||
rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
|
||||
return tuple(_InvalidIndex(*row) for row in rows)
|
||||
|
||||
@staticmethod
|
||||
def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
|
||||
from psycopg import sql
|
||||
|
||||
target: Final = sql.Identifier(index.schema, index.name)
|
||||
if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name):
|
||||
return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover"
|
||||
return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt"
|
||||
|
||||
@staticmethod
|
||||
def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None:
|
||||
import psycopg
|
||||
|
||||
statement, action = ProxyExtrasDBManager._index_repair(index)
|
||||
try:
|
||||
conn.execute(statement)
|
||||
except psycopg.Error as e:
|
||||
logger.warning(
|
||||
"Could not repair invalid index %s.%s, will retry on the next startup. "
|
||||
"If this keeps happening, run `%s` by hand as the index owner. Error: %s",
|
||||
index.schema,
|
||||
index.name,
|
||||
statement.as_string(conn),
|
||||
e,
|
||||
)
|
||||
return
|
||||
logger.info("%s invalid index %s.%s", action, index.schema, index.name)
|
||||
|
||||
@staticmethod
|
||||
def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
|
||||
"""Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
|
||||
INVALID (a migration deadlock between replicas is the usual cause; the
|
||||
retried migration skips them because of IF NOT EXISTS). Never raises:
|
||||
returns True when no invalid index remains, False when the repair was
|
||||
skipped or failed and will be retried on the next startup. Looks in the
|
||||
schema DATABASE_URL names, the only URL Prisma migrates through, but
|
||||
connects over DIRECT_URL when set: the session settings, the advisory
|
||||
lock and REINDEX CONCURRENTLY all need one server session, which a
|
||||
transaction pooler does not give."""
|
||||
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 invalid 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)
|
||||
try:
|
||||
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
|
||||
conn.execute("SET statement_timeout = 0")
|
||||
conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout)))
|
||||
found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
|
||||
if not found:
|
||||
return True
|
||||
logger.warning(
|
||||
"Found %d invalid index(es) left by an interrupted CREATE INDEX "
|
||||
"CONCURRENTLY, rebuilding: %s",
|
||||
len(found),
|
||||
", ".join(f"{index.name} (table size {index.table_size})" for index in found),
|
||||
)
|
||||
lock_row: Final = conn.execute(
|
||||
"SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)
|
||||
).fetchone()
|
||||
if lock_row is None or not lock_row[0]:
|
||||
logger.info("Another replica is already rebuilding the invalid indexes, skipping")
|
||||
return False
|
||||
for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
|
||||
ProxyExtrasDBManager._repair_index(conn, index)
|
||||
remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
|
||||
except psycopg.Error as e:
|
||||
logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
|
||||
return False
|
||||
return not remaining
|
||||
|
||||
@staticmethod
|
||||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
"""
|
||||
|
|
@ -994,6 +1109,7 @@ class ProxyExtrasDBManager:
|
|||
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
|
||||
)
|
||||
if migrated:
|
||||
ProxyExtrasDBManager.repair_invalid_indexes()
|
||||
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
|
||||
return migrated
|
||||
|
||||
|
|
|
|||
264
tests/proxy_migration_tests/test_invalid_index_repair.py
Normal file
264
tests/proxy_migration_tests/test_invalid_index_repair.py
Normal file
|
|
@ -0,0 +1,264 @@
|
|||
import os
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, 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)",
|
||||
)
|
||||
|
||||
HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable"
|
||||
HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx"
|
||||
HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC'
|
||||
LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable"
|
||||
LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx"
|
||||
PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable"
|
||||
PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx"
|
||||
|
||||
|
||||
def _base_url() -> str:
|
||||
return os.environ["DATABASE_URL"].split("?")[0]
|
||||
|
||||
|
||||
def _index_validity(schema: str) -> Mapping[str, bool]:
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT c.relname, i.indisvalid FROM pg_index i "
|
||||
"JOIN pg_class c ON c.oid = i.indexrelid "
|
||||
"JOIN pg_namespace n ON n.oid = c.relnamespace "
|
||||
"WHERE n.nspname = %s",
|
||||
(schema,),
|
||||
).fetchall()
|
||||
return MappingProxyType(dict(rows))
|
||||
|
||||
|
||||
def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None:
|
||||
"""Abort a CONCURRENTLY build while it waits on an older snapshot, the same
|
||||
spot the deadlock loser dies at, so it leaves its index INVALID."""
|
||||
with psycopg.connect(_base_url()) as pin:
|
||||
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
|
||||
pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
|
||||
with psycopg.connect(_base_url(), autocommit=True) as builder:
|
||||
builder.execute("SET statement_timeout = '1s'")
|
||||
with pytest.raises(psycopg.errors.QueryCanceled):
|
||||
builder.execute(statement)
|
||||
|
||||
|
||||
def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None:
|
||||
_interrupt_concurrent_build(
|
||||
schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})'
|
||||
)
|
||||
|
||||
|
||||
def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None:
|
||||
_interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"')
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
|
||||
schema: Final = f"invalid_index_{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}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)'
|
||||
)
|
||||
conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)')
|
||||
|
||||
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 fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
|
||||
"""A brand-new database, what a first deploy sees. A scratch schema would
|
||||
not do: the migrations guard on pg_constraint by name across every schema,
|
||||
so a LiteLLM schema already pushed into public makes them skip and then
|
||||
fail, which is exactly what CI's database looks like."""
|
||||
admin_url: Final = _base_url()
|
||||
name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}"
|
||||
with psycopg.connect(admin_url, autocommit=True) as conn:
|
||||
conn.execute(f'CREATE DATABASE "{name}"')
|
||||
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}")
|
||||
yield "public"
|
||||
|
||||
with psycopg.connect(admin_url, autocommit=True) as conn:
|
||||
conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)')
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None:
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
_leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id")
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False}
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None:
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
_leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX)
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"')
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"')
|
||||
before: Final = _index_validity(scratch_schema)
|
||||
assert len(before) == 4
|
||||
assert set(before.values()) == {False}
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None:
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})')
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None:
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)')
|
||||
conn.execute(
|
||||
f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" '
|
||||
f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)'
|
||||
)
|
||||
conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)')
|
||||
assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
|
||||
assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None:
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
|
||||
with psycopg.connect(_base_url(), autocommit=True) as other_replica:
|
||||
other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,))
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is False
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: False}
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None:
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
|
||||
with psycopg.connect(_base_url()) as pin:
|
||||
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
|
||||
pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"')
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False
|
||||
blocked: Final = _index_validity(scratch_schema)
|
||||
assert blocked[HEALTH_INDEX] is False
|
||||
assert [name for name in blocked if name.endswith("_ccnew")]
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
||||
|
||||
|
||||
def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None:
|
||||
with psycopg.connect(_base_url()) as pin:
|
||||
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
|
||||
pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
|
||||
pinned.set()
|
||||
pin.execute("SELECT pg_sleep(%s)", (seconds,))
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_outlives_a_statement_timeout_passed_through_database_url_options(
|
||||
scratch_schema: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000")
|
||||
pinned: Final = threading.Event()
|
||||
holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0))
|
||||
holder.start()
|
||||
pinned.wait()
|
||||
try:
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
finally:
|
||||
holder.join()
|
||||
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}"
|
||||
index: Final = f"{table}_id_idx"
|
||||
monkeypatch.setenv("DATABASE_URL", _base_url())
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
conn.execute(f'CREATE TABLE public."{table}" (id TEXT)')
|
||||
try:
|
||||
_leave_invalid_index("public", table, index, "id")
|
||||
assert _index_validity("public")[index] is False
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
|
||||
assert _index_validity("public")[index] is True
|
||||
finally:
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
conn.execute(f'DROP TABLE public."{table}"')
|
||||
|
||||
|
||||
def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever")
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is False
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None:
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
with pytest.MonkeyPatch.context() as env:
|
||||
env.setenv("DIRECT_URL", f"{_base_url()}?schema=public")
|
||||
env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}")
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
||||
|
||||
|
||||
def _invalidate_deployed_index(schema: str) -> None:
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"')
|
||||
_leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
|
||||
|
||||
@requires_db
|
||||
@pytest.mark.timeout(300)
|
||||
@pytest.mark.parametrize("use_v2_resolver", [True, False])
|
||||
def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None:
|
||||
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
|
||||
_invalidate_deployed_index(fresh_database)
|
||||
assert _index_validity(fresh_database)[HEALTH_INDEX] is False
|
||||
|
||||
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
|
||||
|
||||
assert _index_validity(fresh_database)[HEALTH_INDEX] is True
|
||||
Loading…
Add table
Reference in a new issue