mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(proxy): coordinate v2 migration startup and add container regression CI
This commit is contained in:
Yuneng Jiang
parent
d565031b9c
commit
94032014df
21 changed files with 1909 additions and 190 deletions
|
|
@ -1,10 +1,32 @@
|
|||
version: 2.1
|
||||
parameters:
|
||||
run_migration_tests:
|
||||
type: boolean
|
||||
default: false
|
||||
migration_candidate_image:
|
||||
type: string
|
||||
default: ""
|
||||
migration_source_sha:
|
||||
type: string
|
||||
default: ""
|
||||
orbs:
|
||||
codecov: codecov/codecov@4.0.1
|
||||
node: circleci/node@5.1.0 # Add this line to declare the node orb
|
||||
win: circleci/windows@5.0 # Add Windows orb
|
||||
|
||||
commands:
|
||||
checkout_migration_source:
|
||||
steps:
|
||||
- run:
|
||||
name: Select the requested migration test revision
|
||||
environment:
|
||||
MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >>
|
||||
command: |
|
||||
if [ -n "$MIGRATION_SOURCE_SHA" ]; then
|
||||
[[ "$MIGRATION_SOURCE_SHA" =~ ^[0-9a-f]{40}$ ]] || exit 1
|
||||
git fetch origin "$MIGRATION_SOURCE_SHA"
|
||||
git checkout --detach "$MIGRATION_SOURCE_SHA"
|
||||
fi
|
||||
skip_if_unrelated_changes:
|
||||
parameters:
|
||||
category:
|
||||
|
|
@ -2853,14 +2875,25 @@ jobs:
|
|||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_migration_source
|
||||
- skip_if_unrelated_changes
|
||||
|
||||
- run:
|
||||
name: Build Docker image
|
||||
environment:
|
||||
MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >>
|
||||
command: |
|
||||
docker build \
|
||||
-t litellm-docker-database:ci \
|
||||
-f docker/Dockerfile.database .
|
||||
if [ -n "$MIGRATION_CANDIDATE_IMAGE" ]; then
|
||||
[[ "$MIGRATION_CANDIDATE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+@sha256:[0-9a-f]{64}$ ]] || exit 1
|
||||
docker pull "$MIGRATION_CANDIDATE_IMAGE"
|
||||
docker tag "$MIGRATION_CANDIDATE_IMAGE" litellm-docker-database:ci
|
||||
else
|
||||
docker build \
|
||||
--label org.opencontainers.image.revision="$(git rev-parse HEAD)" \
|
||||
-t litellm-docker-database:ci \
|
||||
-f docker/Dockerfile.database .
|
||||
fi
|
||||
python3 .circleci/scripts/run_migration_tests.py record-image
|
||||
|
||||
- run:
|
||||
name: Save Docker image to workspace root
|
||||
|
|
@ -2871,6 +2904,79 @@ jobs:
|
|||
root: .
|
||||
paths:
|
||||
- litellm-docker-database.tar.zst
|
||||
- migration-image.json
|
||||
|
||||
migration_startup_tests:
|
||||
parameters:
|
||||
suite:
|
||||
type: enum
|
||||
enum: [startup, recovery, legacy]
|
||||
machine:
|
||||
image: ubuntu-2204:2024.04.1
|
||||
resource_class: large
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
LITELLM_MIGRATION_TESTS: "1"
|
||||
LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci
|
||||
MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres
|
||||
MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres
|
||||
MIGRATION_TEST_OUTPUT: /tmp/migration-results
|
||||
PYTHONPATH: tests/e2e
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_migration_source
|
||||
- install_uv
|
||||
- install_rust
|
||||
- restore_cache:
|
||||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- run:
|
||||
name: Install test dependencies
|
||||
command: uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- run:
|
||||
name: Load the shared candidate and start PostgreSQL
|
||||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker run -d --name migration-postgres \
|
||||
-e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres \
|
||||
-p 5432:5432 \
|
||||
postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Run migration startup regressions
|
||||
environment:
|
||||
MIGRATION_TEST_SUITE: << parameters.suite >>
|
||||
MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >>
|
||||
command: |
|
||||
mkdir -p /tmp/migration-results
|
||||
uv run --no-sync python .circleci/scripts/run_migration_tests.py
|
||||
no_output_timeout: 15m
|
||||
- store_test_results:
|
||||
path: /tmp/migration-results/junit
|
||||
- run:
|
||||
name: Package migration diagnostics
|
||||
when: always
|
||||
command: |
|
||||
mkdir -p /tmp/migration-artifacts
|
||||
if [ -d /tmp/migration-results ]; then
|
||||
tar -czf /tmp/migration-artifacts/diagnostics.tar.gz -C /tmp/migration-results .
|
||||
if [ -f /tmp/migration-results/verdict.json ]; then
|
||||
cp /tmp/migration-results/verdict.json /tmp/migration-artifacts/verdict.json
|
||||
fi
|
||||
fi
|
||||
- store_artifacts:
|
||||
path: /tmp/migration-artifacts
|
||||
destination: migration-results
|
||||
- run:
|
||||
name: Remove migration test containers
|
||||
when: always
|
||||
command: |
|
||||
docker ps -aq --filter label=litellm-migration-test=true | xargs -r docker rm -f
|
||||
docker rm -f migration-postgres || true
|
||||
|
||||
test_bad_database_url:
|
||||
machine:
|
||||
|
|
@ -2915,7 +3021,32 @@ jobs:
|
|||
fi
|
||||
|
||||
workflows:
|
||||
migration_startup:
|
||||
when: << pipeline.parameters.run_migration_tests >>
|
||||
jobs: &migration_jobs
|
||||
- build_docker_database_image
|
||||
- migration_startup_tests:
|
||||
name: migration-startup
|
||||
suite: startup
|
||||
requires: [build_docker_database_image]
|
||||
- migration_startup_tests:
|
||||
name: migration-recovery
|
||||
suite: recovery
|
||||
requires: [build_docker_database_image]
|
||||
- migration_startup_tests:
|
||||
name: migration-legacy-and-pooling
|
||||
suite: legacy
|
||||
requires: [build_docker_database_image]
|
||||
migration_startup_scheduled:
|
||||
triggers:
|
||||
- schedule:
|
||||
cron: "17 0,6,12,18 * * *"
|
||||
filters:
|
||||
branches:
|
||||
only: litellm_internal_staging
|
||||
jobs: *migration_jobs
|
||||
build_and_test:
|
||||
unless: << pipeline.parameters.run_migration_tests >>
|
||||
jobs:
|
||||
- using_litellm_on_windows:
|
||||
filters: &main_branches
|
||||
|
|
|
|||
113
.circleci/scripts/run_migration_tests.py
Normal file
113
.circleci/scripts/run_migration_tests.py
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from xml.etree import ElementTree
|
||||
|
||||
SUITES: Final = {
|
||||
"startup": (("test_startup.py",), 12),
|
||||
"recovery": (("test_recovery.py",), 15),
|
||||
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
|
||||
}
|
||||
|
||||
|
||||
def successful_junit(path: Path, expected: int, exit_code: int) -> bool:
|
||||
if exit_code != 0 or not path.is_file():
|
||||
return False
|
||||
try:
|
||||
root: Final = ElementTree.parse(path).getroot()
|
||||
except ElementTree.ParseError:
|
||||
return False
|
||||
cases: Final = tuple(root.iter("testcase"))
|
||||
identities: Final = frozenset((case.get("classname"), case.get("name")) for case in cases)
|
||||
return len(cases) == len(identities) == expected and all(
|
||||
not any(case.find(tag) is not None for tag in ("failure", "error", "skipped")) for case in cases
|
||||
)
|
||||
|
||||
|
||||
def output(*command: str) -> str:
|
||||
return subprocess.check_output(command, text=True, timeout=90).strip()
|
||||
|
||||
|
||||
def record_image() -> None:
|
||||
source: Final = output("git", "rev-parse", "HEAD")
|
||||
image: Final = output("docker", "image", "inspect", "litellm-docker-database:ci", "--format", "{{.Id}}")
|
||||
revision: Final = output(
|
||||
"docker",
|
||||
"image",
|
||||
"inspect",
|
||||
"litellm-docker-database:ci",
|
||||
"--format",
|
||||
'{{index .Config.Labels "org.opencontainers.image.revision"}}',
|
||||
)
|
||||
assert re.fullmatch(r"[0-9a-f]{40}", source), "Invalid source revision"
|
||||
assert revision == source, "Candidate image revision differs from the tested source"
|
||||
Path("migration-image.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"source_sha": source,
|
||||
"image_id": image,
|
||||
"candidate_image": os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
suite: Final = os.environ["MIGRATION_TEST_SUITE"]
|
||||
files, expected = SUITES[suite]
|
||||
metadata: Final = json.loads(Path("migration-image.json").read_text())
|
||||
assert metadata["source_sha"] == output("git", "rev-parse", "HEAD"), "Image and test source revisions differ"
|
||||
assert metadata["image_id"] == output(
|
||||
"docker", "image", "inspect", os.environ["LITELLM_MIGRATION_TEST_IMAGE"], "--format", "{{.Id}}"
|
||||
), "Loaded image differs from the build output"
|
||||
assert metadata["candidate_image"] == os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""), "Wrong release candidate"
|
||||
destination: Final = Path(os.environ["MIGRATION_TEST_OUTPUT"])
|
||||
junit: Final = destination / "junit" / "results.xml"
|
||||
junit.parent.mkdir(parents=True, exist_ok=True)
|
||||
result: Final = subprocess.run(
|
||||
(
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pytest",
|
||||
*(f"tests/e2e/migrations/{name}" for name in files),
|
||||
"-vv",
|
||||
"--tb=short",
|
||||
"--durations=10",
|
||||
f"--junitxml={junit}",
|
||||
"-o",
|
||||
"addopts=",
|
||||
"--reruns=0",
|
||||
),
|
||||
check=False,
|
||||
timeout=1200,
|
||||
)
|
||||
passed: Final = successful_junit(junit, expected, result.returncode)
|
||||
(destination / "verdict.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
**metadata,
|
||||
"suite": suite,
|
||||
"expected_cases": expected,
|
||||
"passed": passed,
|
||||
"pytest_exit_code": result.returncode,
|
||||
"test_revision": metadata["source_sha"],
|
||||
"workflow_id": os.environ.get("CIRCLE_WORKFLOW_ID", ""),
|
||||
"job_number": os.environ.get("CIRCLE_BUILD_NUM", ""),
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
return 0 if passed else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if sys.argv[1:] == ["record-image"]:
|
||||
record_image()
|
||||
else:
|
||||
raise SystemExit(main())
|
||||
89
litellm-proxy-extras/litellm_proxy_extras/migration_lock.py
Normal file
89
litellm-proxy-extras/litellm_proxy_extras/migration_lock.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
import random
|
||||
import time
|
||||
from collections.abc import Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.prisma_toolchain import MIGRATION_LOCK_TIMEOUT_ENV_VAR, migration_lock_timeout
|
||||
|
||||
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
||||
|
||||
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
|
||||
database_url: Final = environment.get("DATABASE_URL")
|
||||
direct_url: Final = environment.get("DIRECT_URL")
|
||||
if not database_url or not direct_url:
|
||||
return environment
|
||||
schema: Final = next((value for key, value in parse_qsl(urlsplit(database_url).query) if key == "schema"), "public")
|
||||
direct: Final = urlsplit(direct_url)
|
||||
parameters: Final = tuple((key, value) for key, value in parse_qsl(direct.query) if key != "schema")
|
||||
return {
|
||||
**environment,
|
||||
"DATABASE_URL": urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema))))),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LockResult:
|
||||
acquired: bool
|
||||
|
||||
|
||||
def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = MIGRATION_LOCK_KEY) -> bool:
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(_LockResult)) as cursor:
|
||||
row: Final = cursor.execute("SELECT pg_try_advisory_xact_lock(%s) AS acquired", (key,)).fetchone()
|
||||
return row is not None and row.acquired
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MigrationCoordinator:
|
||||
connection: "psycopg.Connection[tuple[object, ...]]"
|
||||
|
||||
def check_connection(self) -> None:
|
||||
self.connection.execute("SELECT 1")
|
||||
|
||||
def acquire_prisma_lock(self) -> None:
|
||||
deadline: Final = time.monotonic() + migration_lock_timeout()
|
||||
while time.monotonic() < deadline:
|
||||
if _try_lock(self.connection, 72707369):
|
||||
return
|
||||
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
|
||||
raise RuntimeError(
|
||||
"Timed out waiting for Prisma's lock to recover migration history. LiteLLM startup has stopped. "
|
||||
"Another migration or a pooled database session may still hold the lock. Check the database lock holder. "
|
||||
"When using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler."
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]:
|
||||
import psycopg
|
||||
|
||||
wait_seconds: Final = migration_lock_timeout()
|
||||
deadline: Final = time.monotonic() + wait_seconds
|
||||
try:
|
||||
with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection:
|
||||
coordinator: Final = MigrationCoordinator(connection)
|
||||
logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
|
||||
while time.monotonic() < deadline:
|
||||
with connection.transaction():
|
||||
if _try_lock(connection):
|
||||
logger.info("Acquired the v2 migration coordinator lock")
|
||||
|
||||
yield coordinator
|
||||
coordinator.check_connection()
|
||||
return
|
||||
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
|
||||
except psycopg.Error as exc:
|
||||
raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc
|
||||
raise RuntimeError(
|
||||
f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
|
||||
f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
|
||||
)
|
||||
158
litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py
Normal file
158
litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
import hashlib
|
||||
import subprocess
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.migration_lock import MigrationCoordinator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MigrationProgress:
|
||||
checksum: str
|
||||
applied_steps_count: int
|
||||
logs: str
|
||||
id: str = ""
|
||||
finished: bool = False
|
||||
|
||||
def confirms_completion(self, script: bytes) -> bool:
|
||||
return (
|
||||
self.applied_steps_count == 1
|
||||
and not self.logs.strip()
|
||||
and self.checksum == hashlib.sha256(script).hexdigest()
|
||||
)
|
||||
|
||||
|
||||
def _migration_records(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, migration: Path
|
||||
) -> tuple[MigrationProgress, ...]:
|
||||
from psycopg import sql
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(MigrationProgress)) as cursor:
|
||||
records: Final = cursor.execute(
|
||||
sql.SQL(
|
||||
"SELECT id, checksum, applied_steps_count, coalesce(logs, '') AS logs, "
|
||||
"finished_at IS NOT NULL AS finished FROM {} "
|
||||
"WHERE migration_name = %s AND rolled_back_at IS NULL"
|
||||
).format(sql.Identifier(schema, "_prisma_migrations")),
|
||||
(migration.parent.name,),
|
||||
).fetchall()
|
||||
return tuple(records)
|
||||
|
||||
|
||||
def recover_completed_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool:
|
||||
"""Finish a proven successful row without erasing its durable completion evidence.
|
||||
|
||||
The caller commits this checkpoint before running another Prisma command.
|
||||
"""
|
||||
from psycopg import sql
|
||||
|
||||
coordinator.acquire_prisma_lock()
|
||||
records: Final = _migration_records(coordinator.connection, schema, migration)
|
||||
unfinished: Final = tuple(record for record in records if not record.finished)
|
||||
script: Final = migration.read_bytes()
|
||||
if not unfinished:
|
||||
return any(record.checksum == hashlib.sha256(script).hexdigest() for record in records)
|
||||
if len(unfinished) != 1 or not unfinished[0].confirms_completion(script):
|
||||
return False
|
||||
progress: Final = unfinished[0]
|
||||
result: Final = coordinator.connection.execute(
|
||||
sql.SQL(
|
||||
"UPDATE {} SET finished_at = current_timestamp "
|
||||
"WHERE id = %s AND checksum = %s AND applied_steps_count = 1 "
|
||||
"AND finished_at IS NULL AND rolled_back_at IS NULL AND coalesce(logs, '') = %s"
|
||||
).format(sql.Identifier(schema, "_prisma_migrations")),
|
||||
(progress.id, progress.checksum, progress.logs),
|
||||
)
|
||||
if result.rowcount != 1:
|
||||
raise RuntimeError("Could not complete the confirmed migration history row; retry startup.")
|
||||
logger.info("Completed migration %s using its successful SQL step and matching checksum", migration.parent.name)
|
||||
return True
|
||||
|
||||
|
||||
def migration_files(directory: Path) -> tuple[tuple[str, str], ...]:
|
||||
return tuple(
|
||||
(path.parent.name, hashlib.sha256(path.read_bytes()).hexdigest())
|
||||
for path in sorted((directory / "migrations").glob("*/migration.sql"))
|
||||
)
|
||||
|
||||
|
||||
def baseline_current_schema(
|
||||
coordinator: MigrationCoordinator,
|
||||
schema: str,
|
||||
migrations_dir: Path,
|
||||
prisma_command: str,
|
||||
prisma_env: Mapping[str, str],
|
||||
) -> None:
|
||||
from psycopg import sql
|
||||
|
||||
packaged_dir: Final = Path(__file__).parent
|
||||
migrations: Final = migration_files(migrations_dir)
|
||||
if (
|
||||
not migrations
|
||||
or migrations != migration_files(packaged_dir)
|
||||
or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes()
|
||||
):
|
||||
raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
|
||||
|
||||
coordinator.acquire_prisma_lock()
|
||||
existing: Final = coordinator.connection.execute(
|
||||
"SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(coordinator.connection),)
|
||||
).fetchone()
|
||||
if existing is not None and existing[0] is not None:
|
||||
return
|
||||
try:
|
||||
prisma_toolchain.run_prisma(
|
||||
(
|
||||
prisma_command,
|
||||
"migrate",
|
||||
"diff",
|
||||
"--from-schema-datasource",
|
||||
str(migrations_dir / "schema.prisma"),
|
||||
"--to-schema-datamodel",
|
||||
str(migrations_dir / "schema.prisma"),
|
||||
"--exit-code",
|
||||
),
|
||||
timeout=prisma_toolchain.prisma_command_timeout(),
|
||||
env=prisma_env,
|
||||
)
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
|
||||
raise RuntimeError(
|
||||
"Cannot automatically baseline this database: its schema has not been verified to match this build. "
|
||||
"Establish the existing migration history before retrying. No schema reconciliation was performed. "
|
||||
"If using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler. "
|
||||
f"Schema verification detail: {exc.stderr}"
|
||||
) from exc
|
||||
|
||||
coordinator.check_connection()
|
||||
ledger: Final = sql.Identifier(schema, "_prisma_migrations")
|
||||
coordinator.connection.execute(
|
||||
sql.SQL(
|
||||
"CREATE TABLE {} (id varchar(36) PRIMARY KEY NOT NULL, checksum varchar(64) NOT NULL, "
|
||||
"finished_at timestamptz, migration_name varchar(255) NOT NULL, logs text, rolled_back_at timestamptz, "
|
||||
"started_at timestamptz NOT NULL DEFAULT now(), applied_steps_count integer NOT NULL DEFAULT 0)"
|
||||
).format(ledger)
|
||||
)
|
||||
with coordinator.connection.cursor() as cursor:
|
||||
cursor.executemany(
|
||||
sql.SQL(
|
||||
"INSERT INTO {} (id, checksum, migration_name, logs, started_at, finished_at) "
|
||||
"VALUES (%s, %s, %s, '', current_timestamp, current_timestamp)"
|
||||
).format(ledger),
|
||||
tuple((str(uuid4()), checksum, name) for name, checksum in migrations),
|
||||
)
|
||||
logger.warning(
|
||||
"Legacy migration history was missing. The existing Prisma schema matches this build; "
|
||||
"adopted %s packaged migrations as a baseline. No schema changes were applied, and "
|
||||
"historical data backfills were not replayed or verified. Continuing startup; "
|
||||
"review any feature-specific backfill requirements.",
|
||||
len(migrations),
|
||||
)
|
||||
|
|
@ -59,6 +59,7 @@ except ImportError:
|
|||
PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
|
||||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
|
||||
MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
|
||||
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
|
||||
|
||||
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
|
||||
|
|
@ -106,6 +107,10 @@ def prisma_command_timeout() -> float:
|
|||
)
|
||||
|
||||
|
||||
def migration_lock_timeout() -> float:
|
||||
return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
|
||||
|
||||
|
||||
def prisma_bootstrap_timeout() -> float:
|
||||
"""Seconds the one-time Node toolchain install may run for."""
|
||||
return _timeout_from_env(
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import shutil
|
|||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
|
|
@ -78,15 +79,10 @@ MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
|
|||
|
||||
@dataclass(frozen=True)
|
||||
class _MigrateAttemptBudget:
|
||||
"""Retries left, and the recoveries already run.
|
||||
|
||||
A recovery that lands something new costs nothing, so a database full of
|
||||
objects `prisma db push` created works through them one per pass. Anything
|
||||
that made no progress spends an attempt, so a stuck run still gives up.
|
||||
"""
|
||||
"""Independent bounds for failed attempts and Prisma lock contention."""
|
||||
|
||||
attempts_left: int
|
||||
recoveries: frozenset[str] = frozenset()
|
||||
contention_seconds_left: float = 600.0
|
||||
|
||||
@property
|
||||
def exhausted(self) -> bool:
|
||||
|
|
@ -99,10 +95,14 @@ class _MigrateAttemptBudget:
|
|||
def spend(self) -> "_MigrateAttemptBudget":
|
||||
return replace(self, attempts_left=self.attempts_left - 1)
|
||||
|
||||
def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget":
|
||||
if recovery in self.recoveries:
|
||||
return self.spend()
|
||||
return replace(self, recoveries=self.recoveries | {recovery})
|
||||
def after_contention(self, elapsed: float) -> "_MigrateAttemptBudget":
|
||||
remaining: Final = self.contention_seconds_left - elapsed
|
||||
if remaining <= 0:
|
||||
raise RuntimeError(
|
||||
"Timed out waiting for Prisma's migration advisory lock. Check the running migration "
|
||||
"or increase LITELLM_MIGRATION_LOCK_TIMEOUT."
|
||||
)
|
||||
return replace(self, contention_seconds_left=remaining)
|
||||
|
||||
|
||||
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
|
||||
|
|
@ -836,12 +836,47 @@ class ProxyExtrasDBManager:
|
|||
|
||||
@staticmethod
|
||||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
if not use_migrate:
|
||||
return ProxyExtrasDBManager._run_database_v2(False)
|
||||
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
|
||||
from litellm_proxy_extras.migration_recovery import baseline_current_schema, recover_completed_migration
|
||||
|
||||
database_url: Final = os.environ.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
raise RuntimeError("DATABASE_URL is required for v2 migrations")
|
||||
lock_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.environ.get("DIRECT_URL") or database_url)
|
||||
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
|
||||
|
||||
def recover_completed(name: str) -> bool:
|
||||
if Path(name).name != name or "\\" in name:
|
||||
return False
|
||||
migration: Final = Path(os.getcwd()) / "migrations" / name / "migration.sql"
|
||||
if not migration.is_file():
|
||||
return False
|
||||
with migration_lock(lock_url) as coordinator:
|
||||
return recover_completed_migration(coordinator, schema, migration)
|
||||
|
||||
def baseline_existing(migrations_dir: str) -> None:
|
||||
with migration_lock(lock_url) as coordinator:
|
||||
baseline_current_schema(
|
||||
coordinator, schema, Path(migrations_dir), _get_prisma_command(), migration_environment(_get_prisma_env())
|
||||
)
|
||||
|
||||
while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
|
||||
continue
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _run_database_v2(
|
||||
use_migrate: bool,
|
||||
recover_completed: Callable[[str], bool] = lambda name: False,
|
||||
baseline_existing: "Callable[[str], None] | None" = None,
|
||||
) -> bool:
|
||||
"""
|
||||
v2 migration resolver (opt-in via --use_v2_migration_resolver).
|
||||
|
||||
Runs `prisma migrate deploy` and handles standard recovery paths
|
||||
(P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a
|
||||
concurrent migrate deploy). Critically, it does
|
||||
Runs `prisma migrate deploy`, baselines verified existing schemas,
|
||||
and recovers confirmed SQL completion or reported deadlocks. It does
|
||||
NOT call `_resolve_all_migrations` — the diff-and-force recovery that
|
||||
caused schema thrashing when two LiteLLM versions contended for the
|
||||
same DB during rolling deploys.
|
||||
|
|
@ -850,10 +885,9 @@ class ProxyExtrasDBManager:
|
|||
is logged as a warning, not a fatal error — users whose DBs got into
|
||||
weird shapes from the old thrashing should still be able to start.
|
||||
|
||||
The retry budget only counts attempts that made no progress: see
|
||||
_MigrateAttemptBudget.
|
||||
False requests a committed recovery checkpoint and another deploy
|
||||
pass. True means every pending migration is complete.
|
||||
"""
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
|
||||
|
||||
if not use_migrate:
|
||||
|
|
@ -886,14 +920,22 @@ class ProxyExtrasDBManager:
|
|||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
deploy_timeout = prisma_migrate_deploy_timeout()
|
||||
budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS)
|
||||
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock_timeout
|
||||
|
||||
migration_env: Final = migration_environment(_get_prisma_env())
|
||||
|
||||
budget = _MigrateAttemptBudget(
|
||||
attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS,
|
||||
contention_seconds_left=migration_lock_timeout(),
|
||||
)
|
||||
try:
|
||||
while not budget.exhausted:
|
||||
attempt_started = time.monotonic()
|
||||
try:
|
||||
result = prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=deploy_timeout,
|
||||
env=_get_prisma_env(),
|
||||
env=migration_env,
|
||||
)
|
||||
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
|
||||
return True
|
||||
|
|
@ -909,8 +951,16 @@ class ProxyExtrasDBManager:
|
|||
next_budget = budget.spend()
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
if "P3005" in (e.stderr or "") and baseline_existing is not None:
|
||||
baseline_existing(migrations_dir)
|
||||
return False
|
||||
failed_migration = ProxyExtrasDBManager._v2_failed_migration_name(e.stderr or "")
|
||||
if failed_migration and recover_completed(failed_migration):
|
||||
return False
|
||||
next_budget = ProxyExtrasDBManager._budget_after_deploy_failure(
|
||||
e, budget, schema_path
|
||||
e,
|
||||
budget,
|
||||
time.monotonic() - attempt_started,
|
||||
)
|
||||
|
||||
if next_budget.attempts_left < budget.attempts_left:
|
||||
|
|
@ -919,19 +969,41 @@ class ProxyExtrasDBManager:
|
|||
|
||||
raise RuntimeError(
|
||||
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
|
||||
"attempts that made no progress (timeouts, deadlock retries, or a "
|
||||
"recovery that had already run once). Check database connectivity, "
|
||||
"attempts that made no progress (timeouts or deadlock retries). Check database connectivity, "
|
||||
"load, and _prisma_migrations ledger state, and raise "
|
||||
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
||||
@staticmethod
|
||||
def _v2_failed_migration_name(stderr: str) -> "str | None":
|
||||
if "P3009" in stderr:
|
||||
match = re.search(r"`(\d+_[^`\r\n]+)`", stderr)
|
||||
return match.group(1) if match else None
|
||||
if "P3018" in stderr:
|
||||
match = re.search(r"Migration name: (\d+_[^\r\n]+)", stderr)
|
||||
return match.group(1) if match else None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
|
||||
from litellm_proxy_extras.migration_lock import migration_environment
|
||||
|
||||
try:
|
||||
prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
|
||||
timeout=prisma_command_timeout(),
|
||||
env=migration_environment(_get_prisma_env()),
|
||||
)
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _budget_after_deploy_failure(
|
||||
error: subprocess.CalledProcessError,
|
||||
budget: "_MigrateAttemptBudget",
|
||||
schema_path: str,
|
||||
attempt_seconds: float = 0.0,
|
||||
) -> "_MigrateAttemptBudget":
|
||||
"""Recover from one failed `prisma migrate deploy`, and price the pass.
|
||||
|
||||
|
|
@ -940,37 +1012,35 @@ class ProxyExtrasDBManager:
|
|||
"""
|
||||
stderr = error.stderr or ""
|
||||
|
||||
if "P3005" in stderr and "database schema is not empty" in stderr:
|
||||
logger.info("Schema exists but no migrations ledger — creating baseline")
|
||||
if ProxyExtrasDBManager._create_baseline_migration(schema_path):
|
||||
return budget.after_recovery("baseline")
|
||||
return budget.spend()
|
||||
|
||||
if "P3009" in stderr:
|
||||
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
|
||||
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
|
||||
name = migration_match.group(1)
|
||||
logger.info(
|
||||
f"Migration {name} failed idempotently — marking applied and retrying"
|
||||
)
|
||||
ProxyExtrasDBManager._mark_migration_applied(name)
|
||||
return budget.after_recovery(f"resolved:{name}")
|
||||
if migration_match:
|
||||
migration_name = migration_match.group(1)
|
||||
migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr)
|
||||
if migration_name:
|
||||
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
|
||||
if ledger_logs is not None and (
|
||||
ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
|
||||
):
|
||||
if ledger_logs and _MIGRATION_DEADLOCK_MARKER in ledger_logs:
|
||||
logger.info(
|
||||
"Migration %s failed in a concurrent migrate deploy "
|
||||
"deadlock race, rolling its ledger row back and retrying",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
return budget.spend()
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
"Migration completion could not be verified. LiteLLM startup has stopped.\n\n"
|
||||
f"Prisma migration history (migration name and start time):\n{stderr}\n\n"
|
||||
"A migration has a start record but no successful completion record. "
|
||||
"LiteLLM cannot determine whether its SQL committed from this record alone. "
|
||||
"Startup stopped to avoid repeating or skipping database changes.\n\n"
|
||||
"Before resolving, stop other migration runners and inspect _prisma_migrations, "
|
||||
"the named migration.sql from this build, database logs, and the actual database objects and data. "
|
||||
"Use the same database and this build's schema and migration files for recovery:\n"
|
||||
"- Only after verifying every migration change is present, run "
|
||||
"prisma migrate resolve --applied <migration_name>, then retry startup.\n"
|
||||
"- Only after verifying no migration changes remain (or fully undoing partial changes), run "
|
||||
"prisma migrate resolve --rolled-back <migration_name>, then retry startup. "
|
||||
"This command updates history; it does not undo SQL.\n"
|
||||
"Replace <migration_name> with the reported name. If the outcome remains uncertain, "
|
||||
"leave migration history unchanged and contact your database administrator. "
|
||||
"Repeated restarts alone will not resolve this state."
|
||||
) from error
|
||||
|
||||
if "P3018" in stderr:
|
||||
|
|
@ -981,25 +1051,13 @@ class ProxyExtrasDBManager:
|
|||
f"and retry.\n\nPrisma error:\n{stderr}"
|
||||
) from error
|
||||
|
||||
migration_match = re.search(r"Migration name: (\d+_\S+)", stderr)
|
||||
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
|
||||
name = migration_match.group(1)
|
||||
migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr)
|
||||
if migration_name and _MIGRATION_DEADLOCK_MARKER in stderr:
|
||||
logger.info(
|
||||
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
|
||||
)
|
||||
ProxyExtrasDBManager._mark_migration_applied(name)
|
||||
return budget.after_recovery(f"resolved:{name}")
|
||||
|
||||
if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
|
||||
logger.info(
|
||||
"Migration %s deadlocked against a concurrent "
|
||||
"migrate deploy, rolling its ledger row back "
|
||||
"and retrying",
|
||||
migration_match.group(1),
|
||||
)
|
||||
ProxyExtrasDBManager._roll_back_migration_best_effort(
|
||||
migration_match.group(1)
|
||||
"Migration %s deadlocked against a concurrent migrate deploy, rolling its ledger row back and retrying",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
return budget.spend()
|
||||
|
||||
raise RuntimeError(
|
||||
|
|
@ -1009,19 +1067,17 @@ class ProxyExtrasDBManager:
|
|||
|
||||
if _MIGRATION_DEADLOCK_MARKER in stderr:
|
||||
logger.info(
|
||||
"prisma migrate deploy attempt %s deadlocked against "
|
||||
"a concurrent migrate deploy, retrying",
|
||||
"prisma migrate deploy attempt %s deadlocked against a concurrent migrate deploy, retrying",
|
||||
budget.attempt_number,
|
||||
)
|
||||
return budget.spend()
|
||||
|
||||
if "P1002" in stderr and "advisory lock" in stderr:
|
||||
logger.info(
|
||||
"prisma migrate deploy attempt %s timed out waiting for "
|
||||
"the advisory lock a concurrent migrate deploy holds, retrying",
|
||||
budget.attempt_number,
|
||||
"Waiting for the advisory lock held by another Prisma migration; "
|
||||
"contention does not spend a migration failure attempt"
|
||||
)
|
||||
return budget.spend()
|
||||
return budget.after_contention(attempt_seconds)
|
||||
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
|
|
|
|||
|
|
@ -32,9 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str):
|
|||
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a permission failure during migrate deploy raises RuntimeError."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
|
|
@ -50,9 +48,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
|||
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
|
|
@ -61,7 +57,7 @@ def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
|||
'Reason: syntax error at or near "BRKN" LINE 42'
|
||||
)
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
|
|
@ -176,51 +172,33 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
|||
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
|
||||
|
||||
|
||||
def test_v2_resolve_specific_migration_failure_raises_runtime_error(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""If marking a migration as applied fails inside P3009 idempotent
|
||||
recovery, the subprocess error must be re-raised as RuntimeError so
|
||||
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
|
||||
def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path):
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "relation already exists")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
ProxyExtrasDBManager,
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: pytest.fail("duplicate-object errors do not prove rollback is safe"),
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("duplicate-object errors do not prove all SQL completed"),
|
||||
)
|
||||
|
||||
# First call: migrate deploy -> P3009 idempotent error.
|
||||
# Recovery path tries _resolve_specific_migration; that also raises.
|
||||
def _failing_resolve(*a, **kw):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd="prisma migrate resolve --applied",
|
||||
stderr="resolve failed",
|
||||
output="",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
|
||||
)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
|
||||
"relation already exists"
|
||||
)
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(
|
||||
RuntimeError, match="Failed to mark migration .* as applied"
|
||||
):
|
||||
stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists"
|
||||
with patch(
|
||||
"litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)
|
||||
) as run:
|
||||
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
|
||||
["migrate", "deploy"],
|
||||
)
|
||||
|
||||
|
||||
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
|
|
@ -252,9 +230,7 @@ _DEADLOCK_P3018_STDERR = (
|
|||
|
||||
def _stub_v2_env(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr("time.sleep", lambda _: None)
|
||||
|
|
@ -272,9 +248,7 @@ def _succeed_after(failures: int, stderr: str):
|
|||
return _OkResult()
|
||||
calls["n"] += 1
|
||||
if calls["n"] <= failures:
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=args[0], stderr=stderr, output=""
|
||||
)
|
||||
raise subprocess.CalledProcessError(returncode=1, cmd=args[0], stderr=stderr, output="")
|
||||
return _OkResult()
|
||||
|
||||
return _run
|
||||
|
|
@ -288,7 +262,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
|
|||
rolled_back = []
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: rolled_back.append(name),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -306,7 +280,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
|
|||
def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
|
||||
"""v2: a deadlock on every attempt still fails after the retry budget."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_v2_roll_back_migration_best_effort", lambda name: None)
|
||||
|
||||
with patch(
|
||||
"litellm_proxy_extras.prisma_toolchain.run_prisma",
|
||||
|
|
@ -335,7 +309,7 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_
|
|||
rolled_back = []
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: rolled_back.append(name),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -350,10 +324,8 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_
|
|||
assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
|
||||
|
||||
|
||||
def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path):
|
||||
"""v2: empty failed ledger logs mean a concurrent deploy moved it on."""
|
||||
def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_path):
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\n"
|
||||
"migrate found failed migrations in the target database\n"
|
||||
|
|
@ -361,22 +333,19 @@ def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path
|
|||
"started at 2026-09-01 18:46:13 UTC failed"
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "")
|
||||
rolled_back = []
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
lambda name: rolled_back.append(name),
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: pytest.fail("empty logs do not prove rollback is safe"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
with patch(
|
||||
"litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)
|
||||
) as run:
|
||||
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
|
||||
["migrate", "deploy"],
|
||||
)
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
|
||||
|
||||
|
||||
def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
|
||||
|
|
@ -392,12 +361,12 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
|
|||
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: pytest.fail("an unreadable ledger must not trigger a retry"),
|
||||
)
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
|
|
@ -418,7 +387,7 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
|
|||
)
|
||||
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ Code-style rules for writing tests under `tests/e2e/`. The harness already encod
|
|||
|
||||
Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family or behavior area. If you add a new folder, you must add a line here describing what kind of tests belong in it, so the layout stays self-describing. `gateway/` is the exception: it holds proxy configuration only and never tests
|
||||
|
||||
- `migrations/` - isolated Docker startup, concurrent migration, crash recovery, and legacy database compatibility. The CircleCI migration workflow enables `LITELLM_MIGRATION_TESTS=1`; these tests own their proxy containers and databases, so they do not use the shared proxy preflight or shared database cleanup
|
||||
|
||||
- `llm_translation/` - LLM endpoint and provider-translation behavior: passthrough, custom pricing, OCR, and the non-chat inference endpoints (`/v1/responses`, `/v1/messages`, `/embeddings`, `/v1/rerank`, `/v1/audio/speech`, `/v1/images/generations`), each against a deployment the test creates via `/model/new` and deletes on teardown
|
||||
- `access_control/` - the gateway's authorization and error-shape contract: per-key model allow-lists, route-group permissions (`allowed_routes`), and unknown-model validation
|
||||
- `embeddings/` - the `/embeddings` endpoint across providers
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient)
|
|||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
config.addinivalue_line("markers", "migration_startup: isolated container startup tests run by the migration CI workflow")
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"e2e: live test that requires a running proxy and real provider keys",
|
||||
|
|
@ -123,6 +124,11 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
|
|||
traffic only after the latency-sensitive suites have finished."""
|
||||
for item in items:
|
||||
attach_result_properties(item)
|
||||
if os.environ.get("LITELLM_MIGRATION_TESTS") != "1":
|
||||
deselected = [item for item in items if item.get_closest_marker("migration_startup") is not None]
|
||||
items[:] = [item for item in items if item.get_closest_marker("migration_startup") is None]
|
||||
if deselected:
|
||||
deselected[0].config.hook.pytest_deselected(items=deselected)
|
||||
items.sort(key=lambda item: item.get_closest_marker("load") is not None)
|
||||
|
||||
|
||||
|
|
@ -155,7 +161,7 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
|
|||
Unmarked tests (unit coverage of the harness) don't touch the proxy, so they
|
||||
run even when none is up. Never skip for a missing proxy. Replay mode needs
|
||||
the proxy too: only provider-bound traffic replays from the bundle."""
|
||||
if item.get_closest_marker("e2e") is None:
|
||||
if item.get_closest_marker("e2e") is None or item.get_closest_marker("migration_startup") is not None:
|
||||
return
|
||||
reason = _proxy_fail_reason()
|
||||
if reason is not None:
|
||||
|
|
@ -168,7 +174,7 @@ def pytest_runtest_call(item: pytest.Item) -> None:
|
|||
guard before truncating the spend-log DB. Tests under `tests/e2e/` without the
|
||||
`e2e` marker (pure unit coverage for the harness itself) never hit the proxy,
|
||||
so they must not arm the destructive DB truncate."""
|
||||
if item.get_closest_marker("e2e") is None:
|
||||
if item.get_closest_marker("e2e") is None or item.get_closest_marker("migration_startup") is not None:
|
||||
return
|
||||
item.session.stash[_E2E_TEST_RAN] = True
|
||||
|
||||
|
|
|
|||
0
tests/e2e/migrations/__init__.py
Normal file
0
tests/e2e/migrations/__init__.py
Normal file
130
tests/e2e/migrations/checks.py
Normal file
130
tests/e2e/migrations/checks.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
import hashlib
|
||||
from contextlib import ExitStack
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
from psycopg import sql
|
||||
|
||||
from .containers import Containers, Replica, failed, until
|
||||
from .database import GATE_KEY, Database
|
||||
from .startup_models import Migration
|
||||
|
||||
COMPLETE_SQL: Final = "CREATE TABLE migration_effect (id int PRIMARY KEY); INSERT INTO migration_effect VALUES (1);"
|
||||
COMPLETE: Final = Migration("20990101000000_startup_test", COMPLETE_SQL)
|
||||
NEXT: Final = Migration(
|
||||
"20990102000000_next_test",
|
||||
"CREATE TABLE migration_next (id int PRIMARY KEY); INSERT INTO migration_next VALUES (2);",
|
||||
)
|
||||
FATAL: Final = Migration(COMPLETE.name, "DO $$ BEGIN RAISE EXCEPTION 'MIGRATION_TEST_FATAL'; END $$;")
|
||||
GATED: Final = Migration(
|
||||
COMPLETE.name, f"SELECT pg_advisory_lock({GATE_KEY}); {COMPLETE.script} SELECT pg_advisory_unlock({GATE_KEY});"
|
||||
)
|
||||
|
||||
|
||||
def start_replicas(
|
||||
stack: ExitStack, containers: Containers, database: Database, migrations: tuple[Migration, ...] = (), count: int = 3
|
||||
) -> tuple[Replica, ...]:
|
||||
return tuple(stack.enter_context(containers.start(database, migrations)) for _ in range(count))
|
||||
|
||||
|
||||
def assert_completed(database: Database, migration: Migration = COMPLETE) -> None:
|
||||
assert database.query(
|
||||
"SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s",
|
||||
(migration.name,),
|
||||
) == ((True, True, 1),), "Expected exactly one successful SQL execution"
|
||||
assert database.query("SELECT id FROM migration_effect") == ((1,),)
|
||||
|
||||
|
||||
def confirmed_history(database: Database) -> str:
|
||||
database.execute(COMPLETE_SQL)
|
||||
row_id: Final = str(uuid4())
|
||||
database.execute(
|
||||
"INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) VALUES (%s, %s, %s, 1)",
|
||||
(row_id, COMPLETE.name, hashlib.sha256(COMPLETE.script.encode()).hexdigest()),
|
||||
)
|
||||
return row_id
|
||||
|
||||
|
||||
def assert_original_proof(database: Database, row_id: str, finished: bool) -> None:
|
||||
assert database.query(
|
||||
"SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM _prisma_migrations WHERE migration_name = %s",
|
||||
(COMPLETE.name,),
|
||||
) == ((row_id, 1, finished, True),), "Recovery lost or replaced the original durable SQL proof"
|
||||
assert database.query("SELECT id FROM migration_effect") == ((1,),)
|
||||
|
||||
|
||||
def pause_completion(database: Database) -> None:
|
||||
database.execute(
|
||||
sql.SQL(
|
||||
"CREATE FUNCTION migration_pause() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN "
|
||||
"IF NEW.migration_name = {name} AND NEW.finished_at IS NOT NULL THEN "
|
||||
"PERFORM pg_advisory_lock({gate}); PERFORM pg_advisory_unlock({gate}); END IF; RETURN NEW; END $$; "
|
||||
"CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW EXECUTE FUNCTION migration_pause()"
|
||||
).format(name=sql.Literal(COMPLETE.name), gate=sql.Literal(GATE_KEY))
|
||||
)
|
||||
|
||||
|
||||
def interrupt_owner(
|
||||
containers: Containers, database: Database, after_commit: bool, *, stop_database_session: bool = True
|
||||
) -> None:
|
||||
if after_commit:
|
||||
pause_completion(database)
|
||||
with database.lock():
|
||||
with containers.start(database, (COMPLETE if after_commit else GATED,)) as owner:
|
||||
until("migration at the intended crash boundary", lambda: bool(database.blocked()))
|
||||
assert database.exists("migration_effect") == after_commit
|
||||
assert database.query(
|
||||
"SELECT finished_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,)
|
||||
) == ((True,),)
|
||||
blocked: Final = database.blocked()
|
||||
assert len(blocked) == 1
|
||||
backend: Final = blocked[0][0]
|
||||
assert owner.state().Running
|
||||
owner.kill()
|
||||
assert owner.state().ExitCode == 137
|
||||
if stop_database_session:
|
||||
database.query("SELECT pg_terminate_backend(%s)", (backend,))
|
||||
until(
|
||||
"terminated migration backend released",
|
||||
lambda: not database.query("SELECT pid FROM pg_stat_activity WHERE pid = %s", (backend,)),
|
||||
)
|
||||
assert database.query(
|
||||
"SELECT finished_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s",
|
||||
(COMPLETE.name,),
|
||||
) == ((True, int(after_commit)),)
|
||||
assert database.exists("migration_effect") == after_commit
|
||||
if not stop_database_session:
|
||||
until(
|
||||
"database backend noticed container death",
|
||||
lambda: not database.query("SELECT pid FROM pg_stat_activity WHERE pid = %s", (backend,)),
|
||||
60,
|
||||
)
|
||||
|
||||
|
||||
def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None:
|
||||
failed(replicas, "Migration completion could not be verified")
|
||||
started: Final = str(
|
||||
database.query(
|
||||
"SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM _prisma_migrations WHERE migration_name = %s",
|
||||
(COMPLETE.name,),
|
||||
)[0][0]
|
||||
)
|
||||
for replica in replicas:
|
||||
assert_guidance(replica.logs(), started)
|
||||
|
||||
|
||||
def assert_guidance(log: str, started: str) -> None:
|
||||
for detail in (
|
||||
COMPLETE.name,
|
||||
started,
|
||||
"cannot determine whether its SQL committed",
|
||||
"_prisma_migrations",
|
||||
"migration.sql",
|
||||
"Only after verifying every migration change is present",
|
||||
"prisma migrate resolve --applied <migration_name>",
|
||||
"Only after verifying no migration changes remain",
|
||||
"prisma migrate resolve --rolled-back <migration_name>",
|
||||
"leave migration history unchanged",
|
||||
"Repeated restarts alone",
|
||||
):
|
||||
assert detail in log, f"Missing recovery guidance: {detail}"
|
||||
62
tests/e2e/migrations/conftest.py
Normal file
62
tests/e2e/migrations/conftest.py
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
from _pytest.fixtures import SubRequest
|
||||
|
||||
from .containers import Containers, docker, ready
|
||||
from .database import Database, Databases
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def migration_image(tmp_path_factory: pytest.TempPathFactory) -> str:
|
||||
configured: Final = os.environ.get("LITELLM_MIGRATION_TEST_IMAGE")
|
||||
assert configured, "LITELLM_MIGRATION_TEST_IMAGE must name the built candidate image"
|
||||
image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}")
|
||||
assert image.startswith("sha256:"), "Unable to identify the candidate image"
|
||||
output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp())))
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
(output / "image.json").write_text(json.dumps({"requested": configured, "image_id": image}))
|
||||
return image
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def databases() -> Databases:
|
||||
admin: Final = os.environ.get("MIGRATION_TEST_ADMIN_URL", "")
|
||||
parsed: Final = urlsplit(admin)
|
||||
assert parsed.hostname in ("127.0.0.1", "localhost"), "Use an isolated loopback PostgreSQL test cluster"
|
||||
assert parsed.port and parsed.path and not parsed.query, "Supply the test cluster port and admin database"
|
||||
container_admin: Final = os.environ.get(
|
||||
"MIGRATION_TEST_CONTAINER_ADMIN_URL",
|
||||
admin.replace("127.0.0.1", "host.docker.internal").replace("localhost", "host.docker.internal"),
|
||||
)
|
||||
return Databases(admin, container_admin)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def migrated_template(
|
||||
databases: Databases, migration_image: str, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[Database]:
|
||||
output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "seed"
|
||||
with databases.create() as database:
|
||||
with Containers(migration_image, output).start(database) as replica:
|
||||
ready((replica,), database)
|
||||
yield database
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def database(databases: Databases, migrated_template: Database) -> Iterator[Database]:
|
||||
with databases.create(migrated_template) as database:
|
||||
yield database
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Containers:
|
||||
configured: Final = os.environ.get("MIGRATION_TEST_OUTPUT")
|
||||
output: Final = Path(configured) / request.node.name if configured else tmp_path
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
return Containers(migration_image, output)
|
||||
203
tests/e2e/migrations/containers.py
Normal file
203
tests/e2e/migrations/containers.py
Normal file
|
|
@ -0,0 +1,203 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import subprocess
|
||||
import time
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
from e2e_http import NoBody, Success, unwrap
|
||||
from models import KeyGenerateBody, KeyGenerateResponse, KeyInfoParams, KeyInfoResponse
|
||||
from transport import HttpTransport
|
||||
|
||||
from .database import Database, prisma_url
|
||||
from .startup_models import ContainerState, Migration, Observation, Readiness
|
||||
|
||||
MASTER_KEY: Final = "sk-migration-ci-fixture"
|
||||
|
||||
|
||||
def docker(*args: str) -> str:
|
||||
result: Final = subprocess.run(("docker", *args), capture_output=True, text=True, timeout=90)
|
||||
assert result.returncode == 0, f"Docker operation failed: {result.stderr}"
|
||||
return result.stdout.strip()
|
||||
|
||||
|
||||
def until(description: str, condition: Callable[[], bool], seconds: float = 150) -> None:
|
||||
deadline: Final = time.monotonic() + seconds
|
||||
while time.monotonic() < deadline:
|
||||
if condition():
|
||||
return
|
||||
time.sleep(0.25)
|
||||
raise AssertionError(f"Timed out waiting for {description}")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Replica:
|
||||
name: str
|
||||
transport: HttpTransport
|
||||
output: Path
|
||||
|
||||
def state(self) -> ContainerState:
|
||||
return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name))
|
||||
|
||||
def observe(self) -> Observation:
|
||||
state: Final = self.state()
|
||||
result: Final = self.transport.get(
|
||||
"/health/readiness", headers=self.transport.master, params=NoBody(), response_type=Readiness, timeout=1
|
||||
)
|
||||
ready: Final = isinstance(result, Success) and result.data.status == "healthy" and result.data.db == "connected"
|
||||
return Observation(None if state.Running else state.ExitCode, ready)
|
||||
|
||||
def logs(self) -> str:
|
||||
result: Final = subprocess.run(("docker", "logs", self.name), capture_output=True, text=True, timeout=30)
|
||||
assert result.returncode == 0, result.stderr
|
||||
return result.stdout + result.stderr
|
||||
|
||||
def kill(self) -> None:
|
||||
if self.state().Running:
|
||||
docker("kill", self.name)
|
||||
|
||||
def usable(self, database: Database) -> None:
|
||||
alias: Final = f"migration-{uuid4().hex}"
|
||||
key: Final = unwrap(
|
||||
self.transport.post(
|
||||
"/key/generate",
|
||||
headers=self.transport.master,
|
||||
json=KeyGenerateBody(key_alias=alias),
|
||||
response_type=KeyGenerateResponse,
|
||||
)
|
||||
).key
|
||||
info: Final = unwrap(
|
||||
self.transport.get(
|
||||
"/key/info",
|
||||
headers=self.transport.master,
|
||||
params=KeyInfoParams(key=key),
|
||||
response_type=KeyInfoResponse,
|
||||
)
|
||||
)
|
||||
assert info.info.key_alias == alias
|
||||
assert database.query(
|
||||
'SELECT key_alias FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
||||
(hashlib.sha256(key.encode()).hexdigest(),),
|
||||
) == ((alias,),)
|
||||
|
||||
|
||||
def ready(replicas: tuple[Replica, ...], database: Database) -> None:
|
||||
def all_ready() -> bool:
|
||||
observations: Final = tuple(replica.observe() for replica in replicas)
|
||||
assert all(item.exit_code is None for item in observations), "Replica exited before readiness"
|
||||
return all(item.ready for item in observations)
|
||||
|
||||
until("every replica ready", all_ready)
|
||||
for replica in replicas:
|
||||
replica.usable(database)
|
||||
|
||||
|
||||
def failed(replicas: tuple[Replica, ...], marker: str) -> None:
|
||||
def all_stopped() -> bool:
|
||||
observations: Final = tuple(replica.observe() for replica in replicas)
|
||||
assert not any(item.ready for item in observations), "Failed migration exposed a ready proxy"
|
||||
return all(item.exit_code is not None for item in observations)
|
||||
|
||||
until("every replica to reject startup", all_stopped)
|
||||
for replica in replicas:
|
||||
assert replica.state().ExitCode != 0, "Failed startup returned success"
|
||||
assert marker in replica.logs(), f"Startup failed outside the expected migration: {marker}"
|
||||
|
||||
|
||||
def waiting(replicas: tuple[Replica, ...], seconds: float) -> None:
|
||||
deadline: Final = time.monotonic() + seconds
|
||||
while time.monotonic() < deadline:
|
||||
assert all(item.exit_code is None and not item.ready for item in (replica.observe() for replica in replicas)), (
|
||||
"Contending replica exited or served early"
|
||||
)
|
||||
time.sleep(0.25)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Containers:
|
||||
image: str
|
||||
output: Path
|
||||
|
||||
@contextmanager
|
||||
def start(
|
||||
self,
|
||||
database: Database,
|
||||
migrations: tuple[Migration, ...] = (),
|
||||
*,
|
||||
v2: bool = True,
|
||||
disabled: bool = False,
|
||||
environment: Mapping[str, str] | None = None,
|
||||
) -> Generator[Replica]:
|
||||
name: Final = f"litellm-migration-{uuid4().hex[:16]}"
|
||||
directory: Final = self.output / name
|
||||
directory.mkdir(parents=True)
|
||||
for migration in migrations:
|
||||
write_migration(directory, migration)
|
||||
(directory / "config.yaml").write_text(
|
||||
"model_list: []\ngeneral_settings:\n master_key: os.environ/LITELLM_MASTER_KEY\n"
|
||||
)
|
||||
env: Final = {
|
||||
"DATABASE_URL": prisma_url(database.container_url, database.schema),
|
||||
"LITELLM_MASTER_KEY": MASTER_KEY,
|
||||
"LITELLM_SALT_KEY": MASTER_KEY,
|
||||
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
|
||||
"LITELLM_TELEMETRY": "False",
|
||||
"LITELLM_LOG": "INFO",
|
||||
"DATABASE_CONNECTION_POOL_LIMIT": "2",
|
||||
"DEFAULT_NUM_WORKERS_LITELLM_PROXY": "1",
|
||||
"USE_V2_MIGRATION_RESOLVER": str(v2).lower(),
|
||||
"DISABLE_SCHEMA_UPDATE": str(disabled).lower(),
|
||||
"LITELLM_MIGRATION_DIR": "/migration-test/prisma",
|
||||
"LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT": "180",
|
||||
**(environment or {}),
|
||||
}
|
||||
try:
|
||||
docker(
|
||||
"run",
|
||||
"-d",
|
||||
"--name",
|
||||
name,
|
||||
"--label",
|
||||
"litellm-migration-test=true",
|
||||
"--add-host",
|
||||
"host.docker.internal:host-gateway",
|
||||
"-p",
|
||||
"127.0.0.1::4000",
|
||||
"-v",
|
||||
f"{directory}:/migration-test",
|
||||
*(arg for key, value in env.items() for arg in ("-e", f"{key}={value}")),
|
||||
self.image,
|
||||
"--config",
|
||||
"/migration-test/config.yaml",
|
||||
"--host",
|
||||
"0.0.0.0",
|
||||
"--port",
|
||||
"4000",
|
||||
)
|
||||
port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1])
|
||||
replica: Final = Replica(name, HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), directory)
|
||||
yield replica
|
||||
finally:
|
||||
try:
|
||||
state: Final = subprocess.run(
|
||||
("docker", "inspect", "--format", "{{json .State}}", name),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
(directory / "state.json").write_text(state.stdout or state.stderr)
|
||||
logs: Final = subprocess.run(("docker", "logs", name), capture_output=True, text=True, timeout=30)
|
||||
(directory / "proxy.log").write_text(logs.stdout + logs.stderr)
|
||||
finally:
|
||||
subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True)
|
||||
|
||||
|
||||
def write_migration(directory: Path, migration: Migration) -> None:
|
||||
path: Final = directory / "prisma" / "migrations" / migration.name
|
||||
path.mkdir(parents=True)
|
||||
(path / "migration.sql").write_text(migration.script)
|
||||
135
tests/e2e/migrations/database.py
Normal file
135
tests/e2e/migrations/database.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, LiteralString
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
from uuid import uuid4
|
||||
|
||||
import psycopg
|
||||
from psycopg import sql
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
Scalar = str | int | bool | None
|
||||
ROWS: Final = TypeAdapter(tuple[tuple[Scalar, ...], ...])
|
||||
GATE_KEY: Final = 39178002
|
||||
PRISMA_LOCK: Final = 72707369
|
||||
COORDINATOR_LOCK: Final = int.from_bytes(b"llm_mig2", "big")
|
||||
|
||||
|
||||
def connect_url(url: str, name: str) -> str:
|
||||
return urlunsplit(urlsplit(url)._replace(path=f"/{name}", query=""))
|
||||
|
||||
|
||||
def prisma_url(url: str, schema: str) -> str:
|
||||
parsed: Final = urlsplit(url)
|
||||
query: Final = tuple((key, value) for key, value in parse_qsl(parsed.query) if key != "schema")
|
||||
return urlunsplit(parsed._replace(query=urlencode((*query, ("schema", schema)))))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Database:
|
||||
name: str
|
||||
url: str
|
||||
container_url: str
|
||||
schema: str = "public"
|
||||
|
||||
@contextmanager
|
||||
def connection(self) -> Generator[psycopg.Connection[tuple[object, ...]]]:
|
||||
with psycopg.connect(self.url, autocommit=True, connect_timeout=5) as connection:
|
||||
connection.execute(sql.SQL("SET search_path TO {}").format(sql.Identifier(self.schema)))
|
||||
connection.execute("SET statement_timeout = '15s'")
|
||||
yield connection
|
||||
|
||||
def execute(self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = ()) -> None:
|
||||
with self.connection() as connection:
|
||||
connection.execute(statement, params or None)
|
||||
|
||||
def query(
|
||||
self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = ()
|
||||
) -> tuple[tuple[Scalar, ...], ...]:
|
||||
with self.connection() as connection:
|
||||
return ROWS.validate_python(connection.execute(statement, params or None).fetchall())
|
||||
|
||||
def exists(self, name: str) -> bool:
|
||||
return self.query("SELECT to_regclass(%s) IS NOT NULL", (name,)) == ((True,),)
|
||||
|
||||
def history(self) -> tuple[tuple[Scalar, ...], ...]:
|
||||
if not self.exists("_prisma_migrations"):
|
||||
return ()
|
||||
return self.query(
|
||||
"SELECT id, migration_name, checksum, started_at::text, finished_at::text, rolled_back_at::text, "
|
||||
"applied_steps_count, logs FROM _prisma_migrations ORDER BY id"
|
||||
)
|
||||
|
||||
def blocked(self, key: int = GATE_KEY) -> tuple[tuple[Scalar, ...], ...]:
|
||||
return self.query(
|
||||
"SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND NOT granted "
|
||||
"AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) "
|
||||
"AND classid = %s AND objid = %s ORDER BY pid",
|
||||
(key >> 32, key & 0xFFFFFFFF),
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def lock(self, key: int = GATE_KEY) -> Generator[None]:
|
||||
with self.connection() as connection:
|
||||
connection.execute("SELECT pg_advisory_lock(%s)", (key,))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
connection.execute("SELECT pg_advisory_unlock(%s)", (key,))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Databases:
|
||||
admin_url: str
|
||||
container_admin_url: str
|
||||
|
||||
@contextmanager
|
||||
def create(self, template: Database | None = None, schema: str = "public") -> Generator[Database]:
|
||||
name: Final = f"litellm_migration_test_{uuid4().hex[:20]}"
|
||||
database: Final = Database(
|
||||
name, connect_url(self.admin_url, name), connect_url(self.container_admin_url, name), schema
|
||||
)
|
||||
with psycopg.connect(self.admin_url, autocommit=True, connect_timeout=5) as connection:
|
||||
connection.execute(
|
||||
sql.SQL("CREATE DATABASE {} TEMPLATE {}").format(
|
||||
sql.Identifier(name), sql.Identifier(template.name if template else "template0")
|
||||
)
|
||||
)
|
||||
try:
|
||||
yield database
|
||||
finally:
|
||||
with psycopg.connect(self.admin_url, autocommit=True, connect_timeout=5) as connection:
|
||||
connection.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name)))
|
||||
|
||||
|
||||
@contextmanager
|
||||
def restricted_user(database: Database) -> Generator[Database]:
|
||||
role: Final = f"migration_reader_{uuid4().hex[:16]}"
|
||||
password: Final = "migration-test-password"
|
||||
with database.connection() as connection:
|
||||
connection.execute(
|
||||
sql.SQL("CREATE ROLE {} LOGIN PASSWORD {}").format(sql.Identifier(role), sql.Literal(password))
|
||||
)
|
||||
try:
|
||||
database.execute(
|
||||
sql.SQL("GRANT USAGE ON SCHEMA {} TO {}").format(sql.Identifier(database.schema), sql.Identifier(role))
|
||||
)
|
||||
database.execute(
|
||||
sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA {} TO {}").format(
|
||||
sql.Identifier(database.schema), sql.Identifier(role)
|
||||
)
|
||||
)
|
||||
local: Final = urlsplit(database.url)
|
||||
remote: Final = urlsplit(database.container_url)
|
||||
yield Database(
|
||||
database.name,
|
||||
urlunsplit(local._replace(netloc=f"{role}:{password}@{local.hostname}:{local.port}")),
|
||||
urlunsplit(remote._replace(netloc=f"{role}:{password}@{remote.hostname}:{remote.port}")),
|
||||
database.schema,
|
||||
)
|
||||
finally:
|
||||
database.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role)))
|
||||
database.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role)))
|
||||
25
tests/e2e/migrations/startup_models.py
Normal file
25
tests/e2e/migrations/startup_models.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class Readiness(BaseModel):
|
||||
status: str = ""
|
||||
db: str = ""
|
||||
|
||||
|
||||
class ContainerState(BaseModel):
|
||||
Running: bool
|
||||
ExitCode: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Observation:
|
||||
exit_code: int | None
|
||||
ready: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Migration:
|
||||
name: str
|
||||
script: str
|
||||
84
tests/e2e/migrations/test_legacy.py
Normal file
84
tests/e2e/migrations/test_legacy.py
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
from contextlib import ExitStack
|
||||
from dataclasses import replace
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
|
||||
from .checks import COMPLETE, assert_completed, confirmed_history, assert_original_proof, start_replicas
|
||||
from .containers import Containers, failed, ready
|
||||
from .database import Database, Databases
|
||||
|
||||
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
|
||||
|
||||
|
||||
def adopt_legacy(containers: Containers, database: Database) -> None:
|
||||
count: Final = database.query("SELECT count(*) FROM _prisma_migrations")[0][0]
|
||||
existing_keys: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')
|
||||
database.execute(
|
||||
"INSERT INTO \"LiteLLM_ShadowEvalJob\" (id, group_id, target_id, router_name, judge_model, shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', 'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())"
|
||||
)
|
||||
database.execute("DROP TABLE _prisma_migrations")
|
||||
with ExitStack() as stack:
|
||||
replicas: Final = start_replicas(stack, containers, database)
|
||||
ready(replicas, database)
|
||||
logs: Final = "\n".join(replica.logs() for replica in replicas)
|
||||
for detail in (
|
||||
"Legacy migration history was missing",
|
||||
"historical data backfills were not replayed or verified",
|
||||
"Continuing startup",
|
||||
):
|
||||
assert detail in logs
|
||||
assert database.query("SELECT count(*) FROM _prisma_migrations") == ((count,),)
|
||||
assert database.query(
|
||||
"SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL OR applied_steps_count <> 0"
|
||||
) == ((0,),)
|
||||
assert set(existing_keys).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token'))
|
||||
assert database.query("SELECT stopped_by FROM \"LiteLLM_ShadowEvalJob\" WHERE id = 'migration-legacy'") == (
|
||||
(None,),
|
||||
)
|
||||
|
||||
|
||||
class TestLegacyMigrations:
|
||||
def test_matching_schema_warns_and_starts(self, containers: Containers, database: Database) -> None:
|
||||
adopt_legacy(containers, database)
|
||||
|
||||
@pytest.mark.parametrize("fault", ("schema_drift", "custom_migrations", "empty_ledger"))
|
||||
def test_unrecognized_legacy_state_is_not_baselined(
|
||||
self, containers: Containers, database: Database, fault: str
|
||||
) -> None:
|
||||
if fault == "empty_ledger":
|
||||
database.execute("TRUNCATE _prisma_migrations")
|
||||
else:
|
||||
database.execute("DROP TABLE _prisma_migrations")
|
||||
if fault == "schema_drift":
|
||||
database.execute('ALTER TABLE "LiteLLM_VerificationToken" DROP COLUMN key_alias CASCADE')
|
||||
with containers.start(database, (COMPLETE,) if fault == "custom_migrations" else ()) as replica:
|
||||
failed((replica,), "Cannot automatically baseline" if fault != "empty_ledger" else "migration")
|
||||
assert not database.exists("migration_effect")
|
||||
if database.exists("_prisma_migrations"):
|
||||
assert database.query(
|
||||
"SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NOT NULL AND applied_steps_count <> 1"
|
||||
) == ((0,),)
|
||||
|
||||
@pytest.mark.parametrize("scenario", ("upgrade", "recovery", "legacy"))
|
||||
def test_non_default_schema(
|
||||
self, containers: Containers, databases: Databases, scenario: Literal["upgrade", "recovery", "legacy"]
|
||||
) -> None:
|
||||
with databases.create(schema="migration tenant") as database:
|
||||
with containers.start(database) as seed:
|
||||
ready((seed,), database)
|
||||
match scenario:
|
||||
case "upgrade":
|
||||
with ExitStack() as stack:
|
||||
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
|
||||
assert_completed(database)
|
||||
case "recovery":
|
||||
original: Final = confirmed_history(database)
|
||||
with ExitStack() as stack:
|
||||
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
|
||||
assert_original_proof(database, original, True)
|
||||
case "legacy":
|
||||
adopt_legacy(containers, database)
|
||||
public: Final = replace(database, schema="public")
|
||||
assert not public.exists("_prisma_migrations")
|
||||
assert not public.exists('"LiteLLM_VerificationToken"')
|
||||
135
tests/e2e/migrations/test_pooling.py
Normal file
135
tests/e2e/migrations/test_pooling.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
import subprocess
|
||||
from collections.abc import Generator
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
from uuid import uuid4
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from psycopg import sql
|
||||
|
||||
from .checks import COMPLETE, assert_completed
|
||||
from .containers import Containers, docker, ready, until
|
||||
from .database import Database, Databases, prisma_url, restricted_user
|
||||
|
||||
POOL_IMAGE: Final = (
|
||||
"ghcr.io/cloudnative-pg/pgbouncer@sha256:e6ddfe22d845e603825e235dd8334b21ecd125abea2a2172478f556b8dee2bb8"
|
||||
)
|
||||
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def application_user(database: Database) -> Generator[Database]:
|
||||
with restricted_user(database) as application:
|
||||
role: Final = sql.Identifier(str(urlsplit(application.url).username))
|
||||
schema: Final = sql.Identifier(database.schema)
|
||||
database.execute(sql.SQL("REVOKE CREATE ON SCHEMA {} FROM PUBLIC").format(schema))
|
||||
for statement in (
|
||||
"GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {} TO {}",
|
||||
"GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA {} TO {}",
|
||||
"ALTER DEFAULT PRIVILEGES IN SCHEMA {} GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO {}",
|
||||
"ALTER DEFAULT PRIVILEGES IN SCHEMA {} GRANT USAGE, SELECT ON SEQUENCES TO {}",
|
||||
):
|
||||
database.execute(sql.SQL(statement).format(schema, role))
|
||||
assert application.query("SELECT has_schema_privilege(current_user, %s, 'CREATE')", (database.schema,)) == (
|
||||
(False,),
|
||||
)
|
||||
yield application
|
||||
|
||||
|
||||
@contextmanager
|
||||
def pool(database: Database, output: Path) -> Generator[str]:
|
||||
name: Final = f"litellm-migration-pool-{uuid4().hex[:12]}"
|
||||
url: Final = urlsplit(database.container_url)
|
||||
directory: Final = output / name
|
||||
directory.mkdir(parents=True)
|
||||
(directory / "users.txt").write_text(f'"{url.username}" "{url.password}"\n')
|
||||
(directory / "pgbouncer.ini").write_text(
|
||||
f"[databases]\n* = host={url.hostname} port={url.port} user={url.username} password={url.password}\n"
|
||||
"[pgbouncer]\nlisten_addr = 0.0.0.0\nlisten_port = 6432\nauth_type = trust\nauth_file = /pool/users.txt\n"
|
||||
"pool_mode = transaction\ndefault_pool_size = 1\nreserve_pool_size = 0\nmax_client_conn = 100\n"
|
||||
"max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = extra_float_digits,options\n"
|
||||
)
|
||||
try:
|
||||
docker(
|
||||
"run",
|
||||
"-d",
|
||||
"--name",
|
||||
name,
|
||||
"--label",
|
||||
"litellm-migration-test=true",
|
||||
"--add-host",
|
||||
"host.docker.internal:host-gateway",
|
||||
"-p",
|
||||
"0.0.0.0::6432",
|
||||
"-v",
|
||||
f"{directory}:/pool:ro",
|
||||
"--entrypoint",
|
||||
"/usr/bin/pgbouncer",
|
||||
POOL_IMAGE,
|
||||
"/pool/pgbouncer.ini",
|
||||
)
|
||||
port: Final = int(docker("port", name, "6432/tcp").splitlines()[0].rsplit(":", 1)[1])
|
||||
local_url: Final = urlunsplit(url._replace(netloc=f"{url.username}:{url.password}@127.0.0.1:{port}"))
|
||||
|
||||
def connected() -> bool:
|
||||
try:
|
||||
with psycopg.connect(local_url, autocommit=True, connect_timeout=2) as connection:
|
||||
return connection.execute("SELECT 1").fetchone() == (1,)
|
||||
except psycopg.Error:
|
||||
return False
|
||||
|
||||
until("PgBouncer ready", connected, 30)
|
||||
yield local_url.replace("127.0.0.1", "host.docker.internal") + "?pgbouncer=true"
|
||||
finally:
|
||||
try:
|
||||
logs: Final = subprocess.run(("docker", "logs", name), text=True, capture_output=True, timeout=30)
|
||||
(directory / "pool.log").write_text(logs.stdout + logs.stderr)
|
||||
finally:
|
||||
subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True)
|
||||
|
||||
|
||||
class TestMigrationPooling:
|
||||
@pytest.mark.parametrize("scenario,replica_count", (("fresh", 3), ("upgrade", 3), ("legacy", 3), ("upgrade", 6)))
|
||||
def test_direct_migrations_with_one_application_backend(
|
||||
self,
|
||||
containers: Containers,
|
||||
databases: Databases,
|
||||
migrated_template: Database,
|
||||
scenario: str,
|
||||
replica_count: int,
|
||||
) -> None:
|
||||
with databases.create(None if scenario == "fresh" else migrated_template) as database:
|
||||
if scenario == "legacy":
|
||||
database.execute("DROP TABLE _prisma_migrations")
|
||||
with (
|
||||
application_user(database) as application,
|
||||
pool(application, containers.output) as pooled_url,
|
||||
ExitStack() as stack,
|
||||
):
|
||||
replicas: Final = tuple(
|
||||
stack.enter_context(
|
||||
containers.start(
|
||||
database,
|
||||
(COMPLETE,) if scenario == "upgrade" else (),
|
||||
environment={
|
||||
"DATABASE_URL": prisma_url(pooled_url, database.schema),
|
||||
"DIRECT_URL": database.container_url,
|
||||
},
|
||||
)
|
||||
)
|
||||
for _ in range(replica_count)
|
||||
)
|
||||
ready(replicas, database)
|
||||
if scenario == "upgrade":
|
||||
assert_completed(database)
|
||||
if scenario == "legacy":
|
||||
assert any(
|
||||
"historical data backfills were not replayed or verified" in replica.logs()
|
||||
for replica in replicas
|
||||
)
|
||||
assert database.query("SELECT count(*) FROM _prisma_migrations WHERE applied_steps_count <> 0") == (
|
||||
(0,),
|
||||
)
|
||||
183
tests/e2e/migrations/test_recovery.py
Normal file
183
tests/e2e/migrations/test_recovery.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
from contextlib import ExitStack
|
||||
from typing import Final, Literal
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from .checks import (
|
||||
COMPLETE,
|
||||
FATAL,
|
||||
GATED,
|
||||
NEXT,
|
||||
assert_completed,
|
||||
confirmed_history,
|
||||
interrupt_owner,
|
||||
assert_original_proof,
|
||||
pause_completion,
|
||||
start_replicas,
|
||||
unconfirmed,
|
||||
)
|
||||
from .containers import Containers, failed, ready, until, waiting
|
||||
from .database import COORDINATOR_LOCK, GATE_KEY, Database
|
||||
from .startup_models import Migration
|
||||
|
||||
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
|
||||
|
||||
|
||||
class TestMigrationRecovery:
|
||||
@pytest.mark.parametrize("after_commit", (False, True))
|
||||
def test_container_owner_crash(self, containers: Containers, database: Database, after_commit: bool) -> None:
|
||||
interrupt_owner(containers, database, after_commit, stop_database_session=False)
|
||||
history: Final = database.history()
|
||||
assert database.query(
|
||||
"SELECT applied_steps_count FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,)
|
||||
) == ((int(after_commit),),)
|
||||
with ExitStack() as stack:
|
||||
successors: Final = start_replicas(stack, containers, database, (COMPLETE if after_commit else GATED,))
|
||||
if after_commit:
|
||||
ready(successors, database)
|
||||
assert_completed(database)
|
||||
else:
|
||||
unconfirmed(successors, database)
|
||||
assert database.history() == history
|
||||
|
||||
@pytest.mark.parametrize("after_commit", (False, True))
|
||||
def test_owner_and_database_session_crash(
|
||||
self, containers: Containers, database: Database, after_commit: bool
|
||||
) -> None:
|
||||
interrupt_owner(containers, database, after_commit)
|
||||
history: Final = database.history()
|
||||
with ExitStack() as stack:
|
||||
successors: Final = start_replicas(stack, containers, database, (COMPLETE,))
|
||||
if after_commit:
|
||||
ready(successors, database)
|
||||
assert_completed(database)
|
||||
return
|
||||
unconfirmed(successors, database)
|
||||
assert database.history() == history
|
||||
with containers.start(database, (COMPLETE,)) as restarted:
|
||||
unconfirmed((restarted,), database)
|
||||
assert database.history() == history
|
||||
|
||||
@pytest.mark.parametrize("later_failure", (False, True))
|
||||
def test_remaining_migrations_after_recovery(
|
||||
self, containers: Containers, database: Database, later_failure: bool
|
||||
) -> None:
|
||||
original: Final = confirmed_history(database)
|
||||
next_migration: Final = Migration(
|
||||
NEXT.name,
|
||||
f"SELECT pg_advisory_lock({GATE_KEY}); "
|
||||
+ (FATAL.script if later_failure else NEXT.script)
|
||||
+ f" SELECT pg_advisory_unlock({GATE_KEY});",
|
||||
)
|
||||
with ExitStack() as stack:
|
||||
with database.lock():
|
||||
owner: Final = stack.enter_context(containers.start(database, (COMPLETE, next_migration)))
|
||||
|
||||
def pending() -> bool:
|
||||
observation: Final = owner.observe()
|
||||
assert observation.exit_code is None and not observation.ready, (
|
||||
"Recovered owner served before pending SQL completed"
|
||||
)
|
||||
return bool(database.blocked())
|
||||
|
||||
until("recovering owner reached the next migration", pending)
|
||||
assert_original_proof(database, original, True)
|
||||
assert not database.exists("migration_next")
|
||||
followers: Final = start_replicas(stack, containers, database, (COMPLETE, next_migration), count=2)
|
||||
replicas: Final = (owner, *followers)
|
||||
waiting(replicas, 1)
|
||||
if later_failure:
|
||||
failed(replicas, NEXT.name)
|
||||
assert database.query(
|
||||
"SELECT finished_at IS NULL, logs LIKE %s FROM _prisma_migrations WHERE migration_name = %s",
|
||||
("%MIGRATION_TEST_FATAL%", NEXT.name),
|
||||
) == ((True, True),)
|
||||
else:
|
||||
ready(replicas, database)
|
||||
assert database.query("SELECT id FROM migration_next") == ((2,),)
|
||||
assert_original_proof(database, original, True)
|
||||
|
||||
def test_second_crash_during_recovery_is_atomic(self, containers: Containers, database: Database) -> None:
|
||||
original: Final = confirmed_history(database)
|
||||
pause_completion(database)
|
||||
with database.lock():
|
||||
with containers.start(database, (COMPLETE,)) as recovering:
|
||||
until("history update blocked before commit", lambda: bool(database.blocked()))
|
||||
assert_original_proof(database, original, False)
|
||||
blocked: Final = database.blocked()
|
||||
assert len(blocked) == 1
|
||||
assert database.query("SELECT pg_terminate_backend(%s)", (blocked[0][0],)) == ((True,),)
|
||||
failed((recovering,), "Lost or could not establish v2 migration coordination")
|
||||
assert_original_proof(database, original, False)
|
||||
with ExitStack() as stack:
|
||||
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
|
||||
assert_original_proof(database, original, True)
|
||||
|
||||
def test_competing_recovery_rechecks_stale_failures(self, containers: Containers, database: Database) -> None:
|
||||
original: Final = confirmed_history(database)
|
||||
with ExitStack() as stack:
|
||||
with database.lock(COORDINATOR_LOCK):
|
||||
replicas: Final = start_replicas(stack, containers, database, (COMPLETE,))
|
||||
until(
|
||||
"all replicas observed the unfinished migration",
|
||||
lambda: all(
|
||||
"Waiting for the v2 migration coordinator lock" in replica.logs() for replica in replicas
|
||||
),
|
||||
)
|
||||
assert_original_proof(database, original, False)
|
||||
ready(replicas, database)
|
||||
assert_original_proof(database, original, True)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fault", ("no_steps", "extra_steps", "failure_logs", "checksum", "duplicate_history", "missing_script")
|
||||
)
|
||||
def test_unproven_history_is_never_repaired(
|
||||
self,
|
||||
containers: Containers,
|
||||
database: Database,
|
||||
fault: Literal["no_steps", "extra_steps", "failure_logs", "checksum", "duplicate_history", "missing_script"],
|
||||
) -> None:
|
||||
confirmed_history(database)
|
||||
match fault:
|
||||
case "no_steps":
|
||||
database.execute(
|
||||
"UPDATE _prisma_migrations SET applied_steps_count = 0 WHERE migration_name = %s", (COMPLETE.name,)
|
||||
)
|
||||
case "extra_steps":
|
||||
database.execute(
|
||||
"UPDATE _prisma_migrations SET applied_steps_count = 2 WHERE migration_name = %s", (COMPLETE.name,)
|
||||
)
|
||||
case "failure_logs":
|
||||
database.execute(
|
||||
"UPDATE _prisma_migrations SET logs = 'permission denied' WHERE migration_name = %s",
|
||||
(COMPLETE.name,),
|
||||
)
|
||||
case "checksum":
|
||||
database.execute(
|
||||
"UPDATE _prisma_migrations SET checksum = %s WHERE migration_name = %s", ("0" * 64, COMPLETE.name)
|
||||
)
|
||||
case "duplicate_history":
|
||||
database.execute(
|
||||
"INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) SELECT %s, migration_name, checksum, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s",
|
||||
(str(uuid4()), COMPLETE.name),
|
||||
)
|
||||
case "missing_script":
|
||||
pass
|
||||
history: Final = database.history()
|
||||
with containers.start(database, () if fault == "missing_script" else (COMPLETE,)) as replica:
|
||||
unconfirmed((replica,), database)
|
||||
assert database.history() == history
|
||||
assert database.query("SELECT id FROM migration_effect") == ((1,),)
|
||||
|
||||
def test_coordinator_timeout_preserves_proof(self, containers: Containers, database: Database) -> None:
|
||||
original: Final = confirmed_history(database)
|
||||
with database.lock(COORDINATOR_LOCK):
|
||||
with containers.start(
|
||||
database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "3"}
|
||||
) as replica:
|
||||
failed((replica,), "Timed out waiting for another v2 migration resolver")
|
||||
assert_original_proof(database, original, False)
|
||||
with containers.start(database, (COMPLETE,)) as replica:
|
||||
ready((replica,), database)
|
||||
assert_original_proof(database, original, True)
|
||||
100
tests/e2e/migrations/test_startup.py
Normal file
100
tests/e2e/migrations/test_startup.py
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
from contextlib import ExitStack
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from .checks import COMPLETE, FATAL, GATED, assert_completed, start_replicas
|
||||
from .containers import Containers, failed, ready, until, waiting
|
||||
from .database import PRISMA_LOCK, Database, Databases, restricted_user
|
||||
from .startup_models import Migration
|
||||
|
||||
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
|
||||
|
||||
|
||||
class TestMigrationStartup:
|
||||
@pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False)))
|
||||
def test_fresh_database(self, containers: Containers, databases: Databases, replicas: int, v2: bool) -> None:
|
||||
with databases.create() as database, ExitStack() as stack:
|
||||
ready(tuple(stack.enter_context(containers.start(database, v2=v2)) for _ in range(replicas)), database)
|
||||
assert database.query(
|
||||
"SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL AND rolled_back_at IS NULL"
|
||||
) == ((0,),)
|
||||
assert database.query("SELECT count(*) > 0 FROM _prisma_migrations") == ((True,),)
|
||||
|
||||
def test_concurrent_upgrade(self, containers: Containers, database: Database) -> None:
|
||||
with ExitStack() as stack:
|
||||
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
|
||||
assert_completed(database)
|
||||
|
||||
def test_waiters_survive_prolonged_contention(self, containers: Containers, database: Database) -> None:
|
||||
with ExitStack() as stack:
|
||||
with database.lock():
|
||||
owner: Final = stack.enter_context(containers.start(database, (GATED,)))
|
||||
until("owner blocked in migration SQL", lambda: bool(database.blocked()))
|
||||
followers: Final = start_replicas(stack, containers, database, (GATED,), count=2)
|
||||
until("both followers attempted Prisma locking", lambda: len(database.blocked(PRISMA_LOCK)) == 2)
|
||||
waiting((owner, *followers), 120)
|
||||
ready((owner, *followers), database)
|
||||
assert_completed(database, GATED)
|
||||
|
||||
def test_lock_deadline_then_restart(self, containers: Containers, database: Database) -> None:
|
||||
history: Final = database.history()
|
||||
with database.lock(PRISMA_LOCK):
|
||||
with containers.start(
|
||||
database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "12"}
|
||||
) as replica:
|
||||
until("Prisma lock contention", lambda: bool(database.blocked(PRISMA_LOCK)))
|
||||
failed((replica,), "Timed out waiting for")
|
||||
assert database.history() == history
|
||||
assert not database.exists("migration_effect")
|
||||
with containers.start(database, (COMPLETE,)) as restarted:
|
||||
ready((restarted,), database)
|
||||
assert_completed(database)
|
||||
|
||||
def test_fatal_sql(self, containers: Containers, database: Database) -> None:
|
||||
with ExitStack() as stack:
|
||||
replicas: Final = start_replicas(stack, containers, database, (FATAL,))
|
||||
failed(replicas, COMPLETE.name)
|
||||
assert database.query(
|
||||
"SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE %s AND finished_at IS NULL",
|
||||
(COMPLETE.name, "%MIGRATION_TEST_FATAL%"),
|
||||
) == ((1,),)
|
||||
|
||||
def test_duplicate_object_does_not_hide_incomplete_sql(self, containers: Containers, database: Database) -> None:
|
||||
database.execute(
|
||||
"CREATE TABLE migration_existing (id int PRIMARY KEY); INSERT INTO migration_existing VALUES (42)"
|
||||
)
|
||||
migration: Final = Migration(
|
||||
COMPLETE.name, "CREATE TABLE migration_existing (id int PRIMARY KEY); " + COMPLETE.script
|
||||
)
|
||||
with ExitStack() as stack:
|
||||
failed(start_replicas(stack, containers, database, (migration,)), COMPLETE.name)
|
||||
assert not database.exists("migration_effect")
|
||||
assert database.query("SELECT id FROM migration_existing") == ((42,),)
|
||||
assert database.query(
|
||||
"SELECT finished_at IS NULL FROM _prisma_migrations WHERE migration_name = %s", (COMPLETE.name,)
|
||||
) == ((True,),)
|
||||
|
||||
@pytest.mark.parametrize("v2", (True, False))
|
||||
def test_restart_preserves_history_and_data(self, containers: Containers, database: Database, v2: bool) -> None:
|
||||
history: Final = database.history()
|
||||
before: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')
|
||||
for _ in range(2):
|
||||
with containers.start(database, v2=v2) as replica:
|
||||
ready((replica,), database)
|
||||
assert database.history() == history
|
||||
assert set(before).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token'))
|
||||
|
||||
def test_disabled_migrations(self, containers: Containers, database: Database) -> None:
|
||||
history: Final = database.history()
|
||||
with containers.start(database, (FATAL,), disabled=True) as replica:
|
||||
ready((replica,), database)
|
||||
assert database.history() == history
|
||||
|
||||
def test_insufficient_privileges(self, containers: Containers, database: Database) -> None:
|
||||
history: Final = database.history()
|
||||
with restricted_user(database) as limited:
|
||||
with containers.start(limited, (COMPLETE,)) as replica:
|
||||
failed((replica,), "permission denied")
|
||||
assert database.history() == history
|
||||
assert not database.exists("migration_effect")
|
||||
|
|
@ -728,12 +728,36 @@ ERROR: relation "SomeTable" already exists
|
|||
"""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pooled,direct,expected",
|
||||
(
|
||||
("postgresql://pool/db?pgbouncer=true", None, "postgresql://pool/db?pgbouncer=true"),
|
||||
("postgresql://pool/db?pgbouncer=true", "postgresql://writer/db", "postgresql://writer/db?schema=public"),
|
||||
(
|
||||
"postgresql://pool/db?schema=tenant%20one&pgbouncer=true",
|
||||
"postgresql://writer/db?sslmode=require&schema=wrong",
|
||||
"postgresql://writer/db?sslmode=require&schema=tenant+one",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_v2_migrations_use_the_direct_connection_with_the_runtime_schema(pooled, direct, expected):
|
||||
from litellm_proxy_extras.migration_lock import migration_environment
|
||||
|
||||
environment = {"DATABASE_URL": pooled, "PRISMA_OFFLINE_MODE": "true"}
|
||||
configured = {**environment, **({"DIRECT_URL": direct} if direct else {})}
|
||||
migrated = migration_environment(configured)
|
||||
|
||||
assert migrated["DATABASE_URL"] == expected
|
||||
assert migrated["PRISMA_OFFLINE_MODE"] == "true"
|
||||
assert configured["DATABASE_URL"] == pooled
|
||||
|
||||
|
||||
class _MigrateDeployHarness:
|
||||
"""Drives _setup_database_v2 with a scripted sequence of
|
||||
`prisma migrate deploy` outcomes, with every recovery command faked out so
|
||||
nothing touches a database or the packaged migrations directory."""
|
||||
|
||||
def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False):
|
||||
def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False, confirmed_migrations=()):
|
||||
import subprocess as subprocess_module
|
||||
|
||||
import litellm_proxy_extras.utils as utils_module
|
||||
|
|
@ -744,16 +768,10 @@ class _MigrateDeployHarness:
|
|||
self._outcomes = list(outcomes)
|
||||
self._repeat_last = repeat_last
|
||||
self._subprocess_module = subprocess_module
|
||||
self.confirmed_migrations = set(confirmed_migrations)
|
||||
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_create_baseline_migration",
|
||||
staticmethod(self._fake_baseline),
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path)))
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
|
|
@ -765,13 +783,15 @@ class _MigrateDeployHarness:
|
|||
staticmethod(self.resolved.append),
|
||||
)
|
||||
monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run)
|
||||
monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {})
|
||||
monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
|
||||
|
||||
self.baseline_succeeds = True
|
||||
|
||||
def _fake_baseline(self, *args, **kwargs):
|
||||
self.baselines += 1
|
||||
return self.baseline_succeeds
|
||||
if not self.baseline_succeeds:
|
||||
raise RuntimeError("The existing schema was not verified")
|
||||
|
||||
def _next_outcome(self):
|
||||
if self._outcomes:
|
||||
|
|
@ -791,79 +811,126 @@ class _MigrateDeployHarness:
|
|||
raise self._subprocess_module.CalledProcessError(1, cmd, stderr=outcome)
|
||||
|
||||
def run(self):
|
||||
return ProxyExtrasDBManager._setup_database_v2(use_migrate=True)
|
||||
while not ProxyExtrasDBManager._run_database_v2(
|
||||
use_migrate=True,
|
||||
recover_completed=self._fake_recovery,
|
||||
baseline_existing=self._fake_baseline,
|
||||
):
|
||||
continue
|
||||
return True
|
||||
|
||||
def _fake_recovery(self, name):
|
||||
if name not in self.confirmed_migrations:
|
||||
return False
|
||||
self.confirmed_migrations.remove(name)
|
||||
self.resolved.append(name)
|
||||
return True
|
||||
|
||||
|
||||
class TestMigrateDeployAttemptAccounting:
|
||||
"""A `prisma db push` database has a full schema and no ledger, so the v2
|
||||
resolver baselines it and then works through every migration whose objects
|
||||
already exist. Those recoveries make progress, so they must not spend the
|
||||
retry budget, which is there to stop a run that is getting nowhere."""
|
||||
|
||||
def test_a_push_created_database_finishes_bootstrapping(
|
||||
self, monkeypatch, tmp_path
|
||||
):
|
||||
already_there = [
|
||||
"20250329084805_new_cron_job_table",
|
||||
"20250806095134_rename_alias_to_server_name_mcp_table",
|
||||
"20260224203854_add_agent_object_permissions_table",
|
||||
"20260301120000_fourth_table",
|
||||
"20260302120000_fifth_table",
|
||||
"20260303120000_sixth_table",
|
||||
]
|
||||
def test_a_push_created_database_finishes_bootstrapping(self, monkeypatch, tmp_path):
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
[_P3005_STDERR]
|
||||
+ [_p3018_stderr(name) for name in already_there]
|
||||
+ ["ok"],
|
||||
[_P3005_STDERR, "ok"],
|
||||
)
|
||||
|
||||
assert harness.run() is True
|
||||
assert harness.baselines == 1
|
||||
assert harness.resolved == already_there
|
||||
assert len(harness.deploy_calls) == len(already_there) + 2
|
||||
assert harness.resolved == []
|
||||
assert len(harness.deploy_calls) == 2
|
||||
|
||||
def test_repeated_recovery_of_one_migration_still_gives_up(
|
||||
self, monkeypatch, tmp_path
|
||||
):
|
||||
def test_repeated_recovery_of_one_migration_still_gives_up(self, monkeypatch, tmp_path):
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
[_p3018_stderr("20250329084805_new_cron_job_table")],
|
||||
repeat_last=True,
|
||||
confirmed_migrations=("20250329084805_new_cron_job_table",),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
harness.run()
|
||||
assert len(harness.deploy_calls) <= _ATTEMPT_BUDGET + 1
|
||||
assert len(harness.deploy_calls) == 2
|
||||
assert harness.resolved == ["20250329084805_new_cron_job_table"]
|
||||
|
||||
def test_timeouts_still_spend_the_budget(self, monkeypatch, tmp_path):
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch, tmp_path, ["timeout"], repeat_last=True
|
||||
)
|
||||
harness = _MigrateDeployHarness(monkeypatch, tmp_path, ["timeout"], repeat_last=True)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
harness.run()
|
||||
assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
|
||||
|
||||
def test_a_baseline_that_never_lands_stops_after_the_budget(
|
||||
self, monkeypatch, tmp_path
|
||||
):
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True
|
||||
)
|
||||
def test_an_unverified_baseline_stops_without_replaying_migrations(self, monkeypatch, tmp_path):
|
||||
harness = _MigrateDeployHarness(monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True)
|
||||
harness.baseline_succeeds = False
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
harness.run()
|
||||
assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
|
||||
assert len(harness.deploy_calls) == 1
|
||||
|
||||
def test_lock_contention_does_not_spend_the_failure_budget(self, monkeypatch, tmp_path):
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
["Error: P1002\nTimed out waiting for the advisory lock"] * 6 + ["ok"],
|
||||
)
|
||||
assert harness.run() is True
|
||||
assert len(harness.deploy_calls) == 7
|
||||
|
||||
def test_duplicate_object_error_without_completion_proof_is_fatal(self, monkeypatch, tmp_path):
|
||||
harness = _MigrateDeployHarness(monkeypatch, tmp_path, [_p3018_stderr("20260101000000_x")])
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
harness.run()
|
||||
assert harness.resolved == []
|
||||
assert len(harness.deploy_calls) == 1
|
||||
|
||||
@pytest.mark.parametrize("name", ("20260101000000_x", "20260101000000_migration with spaces"))
|
||||
def test_an_interrupted_migration_with_confirmed_sql_can_finish(self, monkeypatch, tmp_path, name):
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
[f"Error: P3009\nThe `{name}` migration failed", "ok"],
|
||||
confirmed_migrations=(name,),
|
||||
)
|
||||
assert harness.run() is True
|
||||
assert harness.resolved == [name]
|
||||
|
||||
def test_an_interrupted_migration_without_confirmation_stops(self, monkeypatch, tmp_path):
|
||||
name = "20260101000000_x"
|
||||
started = "2026-09-12 20:15:06.694553 UTC"
|
||||
report = f"Error: P3009\nThe `{name}` migration started at {started} failed"
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
[report],
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="Migration completion could not be verified") as failure:
|
||||
harness.run()
|
||||
message = str(failure.value)
|
||||
assert name in message
|
||||
assert started in message
|
||||
assert "start record but no successful completion record" in message
|
||||
assert "cannot determine whether its SQL committed" in message
|
||||
assert "avoid repeating or skipping database changes" in message
|
||||
assert "_prisma_migrations" in message
|
||||
assert "migration.sql" in message
|
||||
assert "same database" in message
|
||||
assert "Only after verifying every migration change is present" in message
|
||||
assert "prisma migrate resolve --applied <migration_name>" in message
|
||||
assert "Only after verifying no migration changes remain" in message
|
||||
assert "prisma migrate resolve --rolled-back <migration_name>" in message
|
||||
assert "leave migration history unchanged" in message
|
||||
assert "Repeated restarts alone" in message
|
||||
assert report in message
|
||||
assert len(harness.deploy_calls) == 1
|
||||
assert harness.resolved == []
|
||||
|
||||
def test_an_unrecoverable_error_is_not_retried(self, monkeypatch, tmp_path):
|
||||
harness = _MigrateDeployHarness(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
["Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near \"SLECT\"\n"],
|
||||
['Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near "SLECT"\n'],
|
||||
repeat_last=True,
|
||||
)
|
||||
|
||||
|
|
@ -873,6 +940,36 @@ class TestMigrateDeployAttemptAccounting:
|
|||
assert harness.resolved == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"steps,logs,script,expected",
|
||||
(
|
||||
(1, "", b"CREATE TABLE item (id int);", True),
|
||||
(0, "", b"CREATE TABLE item (id int);", False),
|
||||
(0, "already exists", b"CREATE TABLE item (id int);", False),
|
||||
(1, "permission denied", b"CREATE TABLE item (id int);", False),
|
||||
(1, "", b"CREATE TABLE item (id text);", False),
|
||||
(2, "", b"CREATE TABLE item (id int);", False),
|
||||
),
|
||||
)
|
||||
def test_migration_completion_requires_a_matching_successful_script(steps, logs, script, expected):
|
||||
import hashlib
|
||||
|
||||
from litellm_proxy_extras.migration_recovery import MigrationProgress
|
||||
|
||||
progress = MigrationProgress(hashlib.sha256(b"CREATE TABLE item (id int);").hexdigest(), steps, logs)
|
||||
assert progress.confirms_completion(script) is expected
|
||||
|
||||
|
||||
def test_prisma_lock_waiting_has_its_own_deadline():
|
||||
from litellm_proxy_extras.utils import _MigrateAttemptBudget
|
||||
|
||||
budget = _MigrateAttemptBudget(attempts_left=4, contention_seconds_left=2)
|
||||
waiting = budget.after_contention(1)
|
||||
assert waiting.attempts_left == 4
|
||||
with pytest.raises(RuntimeError, match="advisory lock"):
|
||||
waiting.after_contention(2)
|
||||
|
||||
|
||||
class TestJWTKeyMappingCascade:
|
||||
"""Regression tests for issue #33702.
|
||||
|
||||
|
|
|
|||
36
tests/proxy_migration_tests/test_migration_ci.py
Normal file
36
tests/proxy_migration_tests/test_migration_ci.py
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
import importlib.util
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
SCRIPT: Final = Path(__file__).resolve().parents[2] / ".circleci/scripts/run_migration_tests.py"
|
||||
SPEC: Final = importlib.util.spec_from_file_location("migration_ci", SCRIPT)
|
||||
assert SPEC is not None and SPEC.loader is not None
|
||||
MODULE: Final = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(MODULE)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"xml,expected,exit_code,passed",
|
||||
(
|
||||
('<testsuites><testsuite><testcase name="a"/><testcase name="b"/></testsuite></testsuites>', 2, 0, True),
|
||||
('<testsuites><testsuite><testcase name="a"/></testsuite></testsuites>', 2, 0, False),
|
||||
("<testsuite><testcase><failure/></testcase></testsuite>", 1, 0, False),
|
||||
("<testsuite><testcase><error/></testcase></testsuite>", 1, 0, False),
|
||||
("<testsuite><testcase><skipped/></testcase></testsuite>", 1, 0, False),
|
||||
("<testsuite><testcase/></testsuite>", 1, 1, False),
|
||||
("<testsuite><testcase/></testsuite>", 1, 5, False),
|
||||
("<testsuite/>", 1, 0, False),
|
||||
("<broken", 1, 0, False),
|
||||
(None, 1, 0, False),
|
||||
('<testsuite><testcase name="a"/><testcase name="a"/></testsuite>', 2, 0, False),
|
||||
),
|
||||
)
|
||||
def test_only_a_complete_passing_suite_can_certify_an_image(
|
||||
tmp_path: Path, xml: str | None, expected: int, exit_code: int, passed: bool
|
||||
) -> None:
|
||||
path: Final = tmp_path / "results.xml"
|
||||
if xml is not None:
|
||||
path.write_text(xml)
|
||||
assert MODULE.successful_junit(path, expected, exit_code) is passed
|
||||
Loading…
Add table
Reference in a new issue