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

"""Prepare the Node toolchain the Prisma CLI needs, separately from migrations.

The Prisma CLI is a Node program. The first invocation inside a fresh
container installs a private Node runtime and npm-installs the CLI itself,
which can take minutes on a cold or slow machine. Sharing one timeout between
that one-time bootstrap and the migration commands makes a slow bootstrap
indistinguishable from a slow migration, so the bootstrap gets killed long
before it can finish.

A killed bootstrap does not correct itself. The installer leaves its cache
directory behind, and Prisma decides whether to install by testing that
directory for existence alone, so every later attempt skips the install and
then fails on a Node binary that was never written. Deleting a cache directory
that exists without a Node binary is what turns a killed bootstrap back into a
recoverable one.

``prisma migrate deploy`` is the other command whose runtime is not a
constant: it grows with the number of pending migrations, so a fresh database
that has to replay every migration this package ships overruns a per-command
budget sized for the short bookkeeping commands, on a laptop as much as on a
slow CI runner. Migrate deploy therefore runs under its own budget.

The Python ``prisma`` wrapper spawns Node, which spawns the Rust schema
engine, so killing only the wrapper on timeout leaves the engine running with
no parent: it keeps mutating the database after the proxy has given up, holds
Prisma's advisory lock so every retry and every later boot queues behind it,
and dies mid-migration once its pipes close, leaving a half-applied ledger row.
Every Prisma command therefore runs in a process group of its own, and a
timeout kills the whole group.

All three budgets are overridable so an operator can widen them without a
release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install,
``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and
``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. The
per-command budget used to bound migrate deploy as well, so a deployment that
raised it above the deploy default keeps that larger budget for deploy unless
the deploy override says otherwise.
"""

import importlib.util
import math
import os
import shutil
import signal
import subprocess
import sys
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import IO, Optional, Union

from litellm_proxy_extras._logging import logger

try:
    from prisma import config as prisma_config
except ImportError:
    prisma_config = None

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
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0

BOOTSTRAP_ARG = "--version"
PRISMA_CONSOLE_SCRIPT = "prisma"


@dataclass(frozen=True)
class ToolchainBootstrap:
    """Outcome of preparing the Prisma toolchain."""

    healed_incomplete_cache: bool
    ready: bool


def _timeout_from_env(env_var: str, default: float) -> float:
    raw = os.getenv(env_var)
    if raw is None:
        return default
    try:
        seconds = float(raw)
    except ValueError:
        logger.warning(
            "%s=%r is not a number, falling back to %ss", env_var, raw, default
        )
        return default
    if not math.isfinite(seconds) or seconds <= 0:
        logger.warning(
            "%s=%r is not a finite positive number, falling back to %ss",
            env_var,
            raw,
            default,
        )
        return default
    return seconds


def prisma_command_timeout() -> float:
    """Seconds any single Prisma command may run for."""
    return _timeout_from_env(
        PRISMA_COMMAND_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_COMMAND_TIMEOUT
    )


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(
        PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT
    )


def prisma_migrate_deploy_timeout() -> float:
    """Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending."""
    if os.getenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR) is not None:
        return _timeout_from_env(
            PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT
        )
    return max(DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, prisma_command_timeout())


def nodeenv_cache_dir() -> Optional[Path]:
    """Where Prisma installs its private Node runtime, or None if unknowable."""
    override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR)
    if override:
        return Path(override).absolute()
    if prisma_config is not None:
        try:
            return Path(prisma_config.nodeenv_cache_dir).absolute()
        except (OSError, ValueError) as e:
            logger.warning("Could not read the Prisma nodeenv cache dir: %s", e)
    try:
        return Path.home() / ".cache" / "prisma-python" / "nodeenv"
    except RuntimeError:
        logger.warning(
            "No resolvable home directory, cannot locate the Prisma nodeenv cache"
        )
        return None


def node_binary_path(cache_dir: Path) -> Path:
    """Path the Node binary occupies once the toolchain is fully installed."""
    if os.name == "nt":
        return cache_dir / "Scripts" / "node.exe"
    return cache_dir / "bin" / "node"


def heal_incomplete_nodeenv_cache() -> bool:
    """Delete a nodeenv cache directory left without a Node binary.

    Returns True when a half-installed toolchain was removed, so the next
    Prisma invocation reinstalls it instead of failing on a missing binary.
    """
    cache_dir = nodeenv_cache_dir()
    if cache_dir is None:
        return False
    try:
        if not cache_dir.is_dir() or node_binary_path(cache_dir).exists():
            return False
    except OSError as e:
        logger.warning("Could not inspect the Node toolchain at %s: %s", cache_dir, e)
        return False
    logger.warning(
        "Node toolchain at %s has no %s, so a previous install was interrupted. "
        "Removing it so it can be reinstalled.",
        cache_dir,
        node_binary_path(cache_dir).name,
    )
    try:
        shutil.rmtree(cache_dir)
    except OSError as e:
        logger.warning("Could not remove %s: %s", cache_dir, e)
        return False
    return True


def _kill_process_group(process: "subprocess.Popen[str]") -> None:
    if os.name == "nt":
        process.kill()
        return
    try:
        os.killpg(process.pid, signal.SIGKILL)
    except ProcessLookupError:
        return


def prisma_cli_available() -> bool:
    """Whether some way of running the Prisma CLI exists: the console script on PATH or the importable package."""
    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
        return True
    return importlib.util.find_spec(PRISMA_CONSOLE_SCRIPT) is not None


def resolve_prisma_argv(argv: Sequence[str]) -> tuple[str, ...]:
    """Route a bare ``prisma`` command through ``python -m prisma`` when the console script is not on PATH.

    The console script and ``python -m prisma`` are the same entry point, but
    only the module form survives an interpreter whose ``bin`` directory is
    missing from PATH, which is how the proxy gets started under launchers and
    init systems. Any other executable name is left untouched.
    """
    if not argv or argv[0] != PRISMA_CONSOLE_SCRIPT:
        return tuple(argv)
    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
        return tuple(argv)
    return (sys.executable, "-m", PRISMA_CONSOLE_SCRIPT, *argv[1:])


def run_prisma(
    argv: Sequence[str],
    *,
    timeout: float,
    env: Mapping[str, str],
    stdout: Union[IO[str], int, None] = subprocess.PIPE,
    stderr: Optional[int] = subprocess.PIPE,
) -> "subprocess.CompletedProcess[str]":
    """Run one Prisma CLI command in its own process group, bounded by ``timeout``.

    Raises ``subprocess.TimeoutExpired`` once the budget is spent, after killing
    the command together with every process it spawned, and
    ``subprocess.CalledProcessError`` on a non-zero exit. Output is captured as
    text unless ``stdout``/``stderr`` say otherwise.
    """
    with subprocess.Popen(
        resolve_prisma_argv(argv),
        env=env,
        stdout=stdout,
        stderr=stderr,
        text=True,
        start_new_session=True,
    ) as process:
        try:
            out, err = process.communicate(timeout=timeout)
        except BaseException:
            _kill_process_group(process)
            raise
    if process.returncode:
        raise subprocess.CalledProcessError(process.returncode, process.args, out, err)
    return subprocess.CompletedProcess(process.args, process.returncode, out, err)


def ensure_prisma_toolchain(
    prisma_command: str, prisma_env: dict[str, str]
) -> ToolchainBootstrap:
    """Install whatever the Prisma CLI needs to run, under its own timeout.

    Never raises. A toolchain that cannot be prepared is reported so the
    caller can go on and let the real Prisma command produce the real error.
    """
    healed = heal_incomplete_nodeenv_cache()
    timeout = prisma_bootstrap_timeout()
    logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout)
    try:
        run_prisma([prisma_command, BOOTSTRAP_ARG], timeout=timeout, env=prisma_env)
    except subprocess.TimeoutExpired:
        logger.warning(
            "Preparing the Prisma CLI toolchain timed out after %ss. Raise %s "
            "if this machine needs longer to install it.",
            timeout,
            PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR,
        )
        return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
    except subprocess.CalledProcessError as e:
        logger.warning("Preparing the Prisma CLI toolchain failed: %s", e.stderr)
        return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
    except OSError as e:
        logger.warning("Could not run the Prisma CLI: %s", e)
        return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
    logger.info("Prisma CLI toolchain ready")
    return ToolchainBootstrap(healed_incomplete_cache=healed, ready=True)