From 94032014df175c3ec3735ded5144e91b5576547b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 18:24:59 -0700 Subject: [PATCH] fix(proxy): coordinate v2 migration startup and add container regression CI --- .circleci/config.yml | 137 +++++++++++- .circleci/scripts/run_migration_tests.py | 113 ++++++++++ .../litellm_proxy_extras/migration_lock.py | 89 ++++++++ .../migration_recovery.py | 158 ++++++++++++++ .../litellm_proxy_extras/prisma_toolchain.py | 5 + .../litellm_proxy_extras/utils.py | 192 +++++++++++------ .../tests/test_setup_database_fail_fast.py | 109 ++++------ tests/e2e/CLAUDE.md | 2 + tests/e2e/conftest.py | 10 +- tests/e2e/migrations/__init__.py | 0 tests/e2e/migrations/checks.py | 130 +++++++++++ tests/e2e/migrations/conftest.py | 62 ++++++ tests/e2e/migrations/containers.py | 203 ++++++++++++++++++ tests/e2e/migrations/database.py | 135 ++++++++++++ tests/e2e/migrations/startup_models.py | 25 +++ tests/e2e/migrations/test_legacy.py | 84 ++++++++ tests/e2e/migrations/test_pooling.py | 135 ++++++++++++ tests/e2e/migrations/test_recovery.py | 183 ++++++++++++++++ tests/e2e/migrations/test_startup.py | 100 +++++++++ .../test_litellm_proxy_extras_utils.py | 191 ++++++++++++---- .../test_migration_ci.py | 36 ++++ 21 files changed, 1909 insertions(+), 190 deletions(-) create mode 100644 .circleci/scripts/run_migration_tests.py create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_lock.py create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py create mode 100644 tests/e2e/migrations/__init__.py create mode 100644 tests/e2e/migrations/checks.py create mode 100644 tests/e2e/migrations/conftest.py create mode 100644 tests/e2e/migrations/containers.py create mode 100644 tests/e2e/migrations/database.py create mode 100644 tests/e2e/migrations/startup_models.py create mode 100644 tests/e2e/migrations/test_legacy.py create mode 100644 tests/e2e/migrations/test_pooling.py create mode 100644 tests/e2e/migrations/test_recovery.py create mode 100644 tests/e2e/migrations/test_startup.py create mode 100644 tests/proxy_migration_tests/test_migration_ci.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 32d2cf0390c..f6f31651306 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -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 diff --git a/.circleci/scripts/run_migration_tests.py b/.circleci/scripts/run_migration_tests.py new file mode 100644 index 00000000000..56029c406fb --- /dev/null +++ b/.circleci/scripts/run_migration_tests.py @@ -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()) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py new file mode 100644 index 00000000000..e4ccbe585a9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py @@ -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}." + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py new file mode 100644 index 00000000000..9202317c776 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py @@ -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), + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index 9cd48fcf11a..07f83f76d2b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -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( diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2145f891318..2749db5d754 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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 , then retry startup.\n" + "- Only after verifying no migration changes remain (or fully undoing partial changes), run " + "prisma migrate resolve --rolled-back , then retry startup. " + "This command updates history; it does not undo SQL.\n" + "Replace 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. " diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 040d67d25e4..338c571eb4f 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -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) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index a58c13d6a1c..9ef7cacf1a8 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -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 diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 36569896125..4a5f0aa880f 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -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 diff --git a/tests/e2e/migrations/__init__.py b/tests/e2e/migrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py new file mode 100644 index 00000000000..b14631ccad8 --- /dev/null +++ b/tests/e2e/migrations/checks.py @@ -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 ", + "Only after verifying no migration changes remain", + "prisma migrate resolve --rolled-back ", + "leave migration history unchanged", + "Repeated restarts alone", + ): + assert detail in log, f"Missing recovery guidance: {detail}" diff --git a/tests/e2e/migrations/conftest.py b/tests/e2e/migrations/conftest.py new file mode 100644 index 00000000000..735adeedbdb --- /dev/null +++ b/tests/e2e/migrations/conftest.py @@ -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) diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py new file mode 100644 index 00000000000..0f5793b81dd --- /dev/null +++ b/tests/e2e/migrations/containers.py @@ -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) diff --git a/tests/e2e/migrations/database.py b/tests/e2e/migrations/database.py new file mode 100644 index 00000000000..a370c21ba0b --- /dev/null +++ b/tests/e2e/migrations/database.py @@ -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))) diff --git a/tests/e2e/migrations/startup_models.py b/tests/e2e/migrations/startup_models.py new file mode 100644 index 00000000000..03a9b4eda78 --- /dev/null +++ b/tests/e2e/migrations/startup_models.py @@ -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 diff --git a/tests/e2e/migrations/test_legacy.py b/tests/e2e/migrations/test_legacy.py new file mode 100644 index 00000000000..7ba73eb82e0 --- /dev/null +++ b/tests/e2e/migrations/test_legacy.py @@ -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"') diff --git a/tests/e2e/migrations/test_pooling.py b/tests/e2e/migrations/test_pooling.py new file mode 100644 index 00000000000..e0a3693b33e --- /dev/null +++ b/tests/e2e/migrations/test_pooling.py @@ -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,), + ) diff --git a/tests/e2e/migrations/test_recovery.py b/tests/e2e/migrations/test_recovery.py new file mode 100644 index 00000000000..58bf5c348d6 --- /dev/null +++ b/tests/e2e/migrations/test_recovery.py @@ -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) diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py new file mode 100644 index 00000000000..a648218cb26 --- /dev/null +++ b/tests/e2e/migrations/test_startup.py @@ -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") diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index e5826b18668..57133ea95c4 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -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 " in message + assert "Only after verifying no migration changes remain" in message + assert "prisma migrate resolve --rolled-back " 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. diff --git a/tests/proxy_migration_tests/test_migration_ci.py b/tests/proxy_migration_tests/test_migration_ci.py new file mode 100644 index 00000000000..30fe7383c07 --- /dev/null +++ b/tests/proxy_migration_tests/test_migration_ci.py @@ -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", + ( + ('', 2, 0, True), + ('', 2, 0, False), + ("", 1, 0, False), + ("", 1, 0, False), + ("", 1, 0, False), + ("", 1, 1, False), + ("", 1, 5, False), + ("", 1, 0, False), + ("', 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