litellm/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py

import hashlib
import re
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 has_packaged_history(migrations_dir: Path) -> bool:
    """Whether `migrations_dir` carries exactly the migration files and schema.prisma this package ships,
    the only history the boot baselines an existing database onto."""
    packaged_dir: Final = Path(__file__).parent
    migrations: Final = migration_files(migrations_dir)
    return (
        bool(migrations)
        and migrations == migration_files(packaged_dir)
        and (migrations_dir / "schema.prisma").read_bytes() == (packaged_dir / "schema.prisma").read_bytes()
    )


def baseline_current_schema(
    coordinator: MigrationCoordinator,
    schema: str,
    migrations_dir: Path,
    prisma_command: str,
    prisma_env: Mapping[str, str],
) -> None:
    from psycopg import sql

    if not has_packaged_history(migrations_dir):
        raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
    migrations: Final = migration_files(migrations_dir)

    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),
    )


_LINE_COMMENT_RE: Final = re.compile(r"--[^\n]*")
_BLOCK_COMMENT_RE: Final = re.compile(r"/\*.*?\*/", re.DOTALL)
_NO_OP_STATEMENT_RE: Final = re.compile(r"^\s*SELECT\s+1\s*$", re.IGNORECASE)


def is_inert_migration(script: str) -> bool:
    """Whether a migration file changes nothing: only comments and `SELECT 1`, so
    applying it can neither repeat nor skip a database change."""
    stripped: Final = _LINE_COMMENT_RE.sub("", _BLOCK_COMMENT_RE.sub("", script))
    return all(not part.strip() or _NO_OP_STATEMENT_RE.match(part) for part in stripped.split(";"))


def roll_back_failed_inert_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool:
    """Roll back the failed ledger row of a migration whose file in this build is inert,
    so `migrate deploy` applies the inert file on its next pass. The row records an
    earlier build's attempt at SQL this build no longer ships (an index now built by the
    migration job), so no database change can be repeated or skipped by replaying
    the empty file. The caller commits this checkpoint before the next Prisma command.
    """
    from psycopg import sql

    if not is_inert_migration(migration.read_text(encoding="utf-8")):
        return False
    coordinator.acquire_prisma_lock()
    records: Final = _migration_records(coordinator.connection, schema, migration)
    unfinished: Final = tuple(record for record in records if not record.finished)
    if len(unfinished) != 1:
        return False
    result: Final = coordinator.connection.execute(
        sql.SQL(
            "UPDATE {} SET rolled_back_at = current_timestamp "
            "WHERE id = %s AND finished_at IS NULL AND rolled_back_at IS NULL"
        ).format(sql.Identifier(schema, "_prisma_migrations")),
        (unfinished[0].id,),
    )
    if result.rowcount != 1:
        raise RuntimeError("Could not roll back the failed inert migration history row; rerun the database setup.")
    logger.info(
        "Rolled back the failed history row of %s: this build ships it as an inert migration, "
        "its index is built by the migration job",
        migration.parent.name,
    )
    return True