litellm/tests/unit/proxy/db/test_prisma_client.py

import asyncio
import importlib.resources
import json
import os
import signal
import sys
import urllib.parse
from typing import Any, Final, cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch

import pytest
from fastapi.testclient import TestClient



from litellm.proxy.db.prisma_client import (
    PrismaManager,
    PrismaWrapper,
    _PrismaDrainTracker,
    _TrackedPrismaEngine,
    should_update_prisma_schema,
)


@pytest.fixture(autouse=True)
def mock_prisma_binary():
    """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests."""
    mock_module = MagicMock()
    with patch.dict(sys.modules, {"prisma": mock_module}):
        yield mock_module


def test_should_update_prisma_schema(monkeypatch):
    # CASE 1: Environment variable behavior
    # When DISABLE_SCHEMA_UPDATE is not set -> should update
    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
    assert should_update_prisma_schema() == True

    # When DISABLE_SCHEMA_UPDATE="true" -> should not update
    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "true")
    assert should_update_prisma_schema() == False

    # When DISABLE_SCHEMA_UPDATE="false" -> should update
    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "false")
    assert should_update_prisma_schema() == True

    # CASE 2: Explicit parameter behavior (overrides env var)
    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
    assert should_update_prisma_schema(True) == False  # Param True -> should not update

    monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)  # Set env var opposite to param
    assert should_update_prisma_schema(False) == True  # Param False -> should update


@pytest.mark.asyncio
async def test_recreate_prisma_client_successful_disconnect():
    """
    Test that recreate_prisma_client works normally when disconnect succeeds.
    """
    # Mock the original prisma client
    mock_prisma = AsyncMock()

    # Create a mock PrismaWrapper instance
    wrapper = Mock()
    wrapper._original_prisma = mock_prisma

    # Configure disconnect to succeed
    mock_prisma.disconnect.return_value = None

    # Mock the entire recreate_prisma_client method to avoid import issues
    async def mock_recreate_prisma_client(new_db_url: str, http_client=None):
        try:
            await mock_prisma.disconnect()
        except Exception:
            pass

        mock_new_prisma = AsyncMock()
        wrapper._original_prisma = mock_new_prisma
        await mock_new_prisma.connect()

    # Assign the mock method to the wrapper
    wrapper.recreate_prisma_client = mock_recreate_prisma_client

    # Call the method
    await wrapper.recreate_prisma_client("postgresql://new:new@localhost:5432/new")

    # Verify that disconnect was called
    mock_prisma.disconnect.assert_called_once()

    # Verify that the new client replaced the original
    assert wrapper._original_prisma != mock_prisma
    assert hasattr(wrapper._original_prisma, "connect")


@pytest.mark.asyncio
async def test_recreate_prisma_client_kills_old_engine_on_disconnect_failure(
    mock_prisma_binary,
):
    """When disconnect() fails, recreate_prisma_client must SIGTERM/SIGKILL the old engine PID."""
    mock_prisma = AsyncMock()
    mock_prisma.disconnect.side_effect = Exception("engine hung")
    mock_prisma.is_connected = MagicMock(return_value=True)

    # Simulate engine subprocess with a known PID
    mock_engine = MagicMock()
    mock_engine.process.pid = 12345
    mock_prisma._engine = mock_engine

    wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)

    # Configure the mock Prisma constructor
    mock_new_prisma = AsyncMock()
    mock_prisma_binary.Prisma.return_value = mock_new_prisma

    with (
        patch("os.kill") as mock_kill,
        patch("asyncio.sleep", new_callable=AsyncMock),
    ):
        await wrapper.recreate_prisma_client("postgresql://new")

    # Verify old engine was killed
    mock_kill.assert_any_call(12345, signal.SIGTERM)
    # Verify new client was created and connected
    mock_new_prisma.connect.assert_awaited_once()


@pytest.mark.asyncio
async def test_recreate_prisma_client_skips_kill_on_successful_disconnect(
    mock_prisma_binary,
):
    """When disconnect() succeeds, no kill should be attempted."""
    mock_prisma = AsyncMock()
    mock_prisma.is_connected = MagicMock(return_value=True)
    mock_prisma.disconnect.return_value = None

    wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)

    mock_new_prisma = AsyncMock()
    mock_prisma_binary.Prisma.return_value = mock_new_prisma

    with patch("os.kill") as mock_kill:
        await wrapper.recreate_prisma_client("postgresql://new")

    mock_kill.assert_not_called()
    mock_new_prisma.connect.assert_awaited_once()


@pytest.mark.asyncio
async def test_recreate_prisma_client_handles_missing_engine_pid(
    mock_prisma_binary,
):
    """When engine PID is unavailable (no _engine attr), kill is skipped gracefully."""
    mock_prisma = AsyncMock()
    mock_prisma.is_connected = MagicMock(return_value=True)
    mock_prisma.disconnect.side_effect = Exception("engine hung")
    mock_prisma._engine = None  # No engine subprocess

    wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)

    mock_new_prisma = AsyncMock()
    mock_prisma_binary.Prisma.return_value = mock_new_prisma

    with (
        patch("os.kill") as mock_kill,
        patch("asyncio.sleep", new_callable=AsyncMock),
    ):
        await wrapper.recreate_prisma_client("postgresql://new")

    mock_kill.assert_not_called()  # PID was 0, kill skipped
    mock_new_prisma.connect.assert_awaited_once()


def test_get_engine_pid_returns_zero_for_disconnected_client(disconnected_prisma):
    """A disconnected client must read as "no engine" instead of raising,
    otherwise the reconnect path can never recover."""
    wrapper = PrismaWrapper(
        original_prisma=disconnected_prisma, iam_token_db_auth=False
    )

    assert wrapper._get_engine_pid() == 0


@pytest.mark.asyncio
async def test_recreate_prisma_client_recovers_from_disconnected_client(
    mock_prisma_binary, disconnected_prisma
):
    """recreate_prisma_client must still build a replacement client when the
    current one is disconnected."""
    wrapper = PrismaWrapper(
        original_prisma=disconnected_prisma, iam_token_db_auth=False
    )

    mock_new_prisma = AsyncMock()
    mock_prisma_binary.Prisma.return_value = mock_new_prisma

    with patch("os.kill") as mock_kill:
        result = await wrapper.recreate_prisma_client("postgresql://new")

    assert result is True
    mock_kill.assert_not_called()
    assert wrapper._original_prisma is mock_new_prisma
    mock_new_prisma.connect.assert_awaited_once()


DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"]
ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV = [
    "db",
    "execute",
    "--schema",
    "schema.prisma",
    "--file",
    str(
        importlib.resources.files("litellm_proxy_extras")
        / "migrations"
        / "20261015120000_add_environment_variables_revision_trigger"
        / "migration.sql"
    ),
]


def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url):
    """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call
    into the opt-in REPLICA IDENTITY FULL step."""
    from litellm.proxy.db.prisma_client import PrismaManager
    from litellm_proxy_extras.replica_identity import REPLICA_IDENTITY_FULL_ENV_VAR
    from litellm_proxy_extras.utils import ProxyExtrasDBManager

    monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
    applied = []
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "apply_replica_identity_full_if_requested",
        staticmethod(lambda: applied.append(True)),
    )

    assert PrismaManager.setup_database(use_migrate=False) is True

    assert fake_prisma_cli.calls == [DB_PUSH_ARGV, ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV]
    assert applied == [True]


def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
    """A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the
    primary key back to ("request_id"), which Postgres rejects; the guard must
    fail fast with guidance instead of running the push."""
    from litellm.proxy.db.prisma_client import PrismaManager
    from litellm_proxy_extras.utils import (
        PARTITIONED_SPEND_LOGS_PUSH_ERROR,
        ProxyExtrasDBManager,
    )

    monkeypatch.setattr(
        ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
    )
    with pytest.raises(RuntimeError) as err:
        PrismaManager.setup_database(use_migrate=False)

    assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
    assert fake_prisma_cli.calls == []


def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
    from litellm.proxy.db.prisma_client import PrismaManager
    from litellm_proxy_extras.utils import ProxyExtrasDBManager

    monkeypatch.setattr(
        ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False)
    )
    assert PrismaManager.setup_database(use_migrate=False) is True

    assert fake_prisma_cli.calls == [DB_PUSH_ARGV, ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV]


def _entra_jwt(expires_in_seconds: int) -> str:
    """A JWT shaped like a real Entra access token, expiring ``expires_in_seconds`` from now."""
    import base64
    from datetime import datetime, timedelta, timezone

    exp = int((datetime.now(tz=timezone.utc) + timedelta(seconds=expires_in_seconds)).timestamp())
    payload = base64.urlsafe_b64encode(json.dumps({"exp": exp}).encode()).rstrip(b"=").decode()
    return f"aGVhZGVy.{payload}.c2ln"


@pytest.fixture
def azure_env(monkeypatch, unset_database_url):
    monkeypatch.setenv("DATABASE_HOST", "pg.postgres.database.azure.com")
    monkeypatch.setenv("DATABASE_PORT", "5432")
    monkeypatch.setenv("DATABASE_USER", "litellm@contoso.onmicrosoft.com")
    monkeypatch.setenv("DATABASE_NAME", "litellm_db")


def _azure_wrapper(token: str, **kwargs):
    from litellm.proxy.db.token_auth import AzureEntraTokenAuth

    return PrismaWrapper(
        original_prisma=MagicMock(),
        token_auth=AzureEntraTokenAuth(token_provider=lambda: token),
        **kwargs,
    )


def test_azure_entra_mint_writes_an_encoded_url_into_the_db_url_env_var(azure_env):
    """The UPN user and the JWT both have to survive being embedded in a URL."""
    token = _entra_jwt(3600)
    wrapper = _azure_wrapper(token)

    db_url = wrapper.get_rds_iam_token()

    assert db_url == (
        f"postgresql://litellm%40contoso.onmicrosoft.com:{urllib.parse.quote(token, safe='')}"
        "@pg.postgres.database.azure.com:5432/litellm_db"
    )
    assert os.environ["DATABASE_URL"] == db_url


@pytest.mark.parametrize(
    ("previous_query", "expected_query"),
    [
        ("max_idle_connection_lifetime=60", {"max_idle_connection_lifetime": ["60"]}),
        (
            "connection_limit=20&pgbouncer=true&max_idle_connection_lifetime=45",
            {"connection_limit": ["20"], "pgbouncer": ["true"], "max_idle_connection_lifetime": ["45"]},
        ),
        (
            "sslmode=require&sslcert=/certs/root.pem&sslaccept=strict&schema=tenant",
            {"sslmode": ["require"], "sslcert": ["/certs/root.pem"], "sslaccept": ["strict"]},
        ),
    ],
)
def test_token_refresh_keeps_the_connection_params_of_the_url_it_replaces(
    azure_env, monkeypatch, previous_query, expected_query
):
    old_token = _entra_jwt(60)
    monkeypatch.setenv(
        "DATABASE_URL",
        f"postgresql://litellm%40contoso.onmicrosoft.com:{urllib.parse.quote(old_token, safe='')}"
        f"@pg.postgres.database.azure.com:5432/litellm_db?{previous_query}",
    )
    new_token = _entra_jwt(3600)

    db_url = _azure_wrapper(new_token).get_rds_iam_token()

    assert db_url is not None
    assert os.environ["DATABASE_URL"] == db_url
    assert urllib.parse.quote(new_token, safe="") in db_url
    assert urllib.parse.parse_qs(urllib.parse.urlsplit(db_url).query) == expected_query


def test_token_refresh_keeps_the_reader_url_params_separate_from_the_writer(azure_env, monkeypatch):
    from litellm.proxy.db.token_auth import IAMEndpoint

    monkeypatch.setenv("DATABASE_URL", "postgresql://w:t@pg:5432/litellm_db?max_idle_connection_lifetime=45")
    monkeypatch.setenv(
        "DATABASE_URL_READ_REPLICA", "postgresql://r:t@replica:5432/litellm_db?max_idle_connection_lifetime=60"
    )
    reader = _azure_wrapper(
        _entra_jwt(3600),
        db_url_env_var="DATABASE_URL_READ_REPLICA",
        iam_endpoint=IAMEndpoint(host="replica", port="5432", user="r", name="litellm_db", schema=None),
    )

    reader_url = reader.get_rds_iam_token()

    assert reader_url is not None
    assert reader_url.startswith("postgresql://r:")
    assert urllib.parse.parse_qs(urllib.parse.urlsplit(reader_url).query) == {"max_idle_connection_lifetime": ["60"]}
    assert os.environ["DATABASE_URL"].endswith("?max_idle_connection_lifetime=45")


def test_azure_entra_refresh_is_scheduled_off_the_jwt_expiry(azure_env):
    """Without reading `exp` this falls back to a fixed 600s interval, which silently
    outlives a token and breaks every reconnect after it lapses (issue #29661)."""
    wrapper = _azure_wrapper(_entra_jwt(3600))
    wrapper.get_rds_iam_token()

    seconds = wrapper._calculate_seconds_until_refresh()

    expected = 3600 - PrismaWrapper.TOKEN_REFRESH_BUFFER_SECONDS
    assert seconds != PrismaWrapper.FALLBACK_REFRESH_INTERVAL_SECONDS
    assert expected - 5 <= seconds <= expected


def test_a_token_whose_expiry_never_advances_cannot_spin_the_refresh_loop(azure_env):
    """azure-identity hands back its cached token when a renewal attempt fails inside its
    own window, so a transient Entra or IMDS problem in the last 3 minutes of a token
    yields a successful refresh whose `exp` has not moved. With no floor on the sleep the
    loop then re-mints and recreates the query engine on every pass, with nothing in
    between, for as long as Entra stays sick."""
    wrapper = _azure_wrapper(_entra_jwt(60))
    wrapper.get_rds_iam_token()
    first = wrapper._calculate_seconds_until_refresh()

    wrapper.get_rds_iam_token()
    second = wrapper._calculate_seconds_until_refresh()

    assert first == second == PrismaWrapper.TOKEN_REFRESH_MIN_SLEEP_SECONDS


def test_azure_entra_token_expiry_is_detected(azure_env):
    wrapper = _azure_wrapper(_entra_jwt(3600))
    fresh_url = wrapper.get_rds_iam_token()
    expired_url = _azure_wrapper(_entra_jwt(-1)).get_rds_iam_token()

    assert wrapper.is_token_expired(fresh_url) is False
    assert wrapper.is_token_expired(expired_url) is True


@pytest.mark.asyncio
async def test_azure_entra_strategy_starts_the_refresh_task(azure_env):
    """The refresh loop is gated on the legacy boolean, so an Azure strategy has to
    get past that gate; a password-auth wrapper still must not start a task."""
    wrapper = _azure_wrapper(_entra_jwt(3600))
    wrapper.get_rds_iam_token()
    password_wrapper = PrismaWrapper(original_prisma=MagicMock())

    await wrapper.start_token_refresh_task()
    await password_wrapper.start_token_refresh_task()
    try:
        assert wrapper._token_refresh_task is not None
        assert not wrapper._token_refresh_task.done()
        assert password_wrapper._token_refresh_task is None
    finally:
        await wrapper.stop_token_refresh_task()


def test_azure_entra_strategy_reads_as_token_auth_enabled(azure_env):
    """`routing_prisma_wrapper` gates the reader's refresh on this flag, so an Azure
    reader has to answer True to it."""
    wrapper = _azure_wrapper(_entra_jwt(3600))

    assert wrapper.iam_token_db_auth is True
    assert wrapper.token_label == "Azure Entra token"


def test_the_token_strategy_cannot_be_swapped_after_construction(azure_env):
    """Assigning the legacy boolean used to replace a configured Entra strategy with the
    RDS one, which points boto at an Azure host."""
    wrapper = _azure_wrapper(_entra_jwt(3600))

    with pytest.raises(AttributeError):
        wrapper.iam_token_db_auth = True


def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch):
    """A blank host used to produce `postgresql://:<token>@:5432/`, which fails deep
    inside Prisma instead of at the misconfiguration."""
    monkeypatch.delenv("DATABASE_HOST")
    wrapper = _azure_wrapper(_entra_jwt(3600))

    with pytest.raises(RuntimeError, match="DATABASE_HOST"):
        wrapper.get_rds_iam_token()


@pytest.mark.timeout(45)
def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch):
    """
    A timed-out `db push` used to leave Node and the schema engine writing the schema,
    so the next attempt pushed into a database the abandoned one was still mutating.
    """
    monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False)
    monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")

    assert PrismaManager.setup_database(use_migrate=False) is True
    assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV, ENVIRONMENT_VARIABLES_REVISION_TRIGGER_ARGV]
    assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)


def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot(
    fake_prisma_cli, unset_database_url, monkeypatch
):
    """
    An ImportError out of setup_database escapes the caller's RuntimeError handler and
    kills boot with a traceback instead of the failed-setup message and exit code.
    """
    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)

    assert PrismaManager.setup_database(use_migrate=False) is False
    assert fake_prisma_cli.calls == []


@pytest.mark.parametrize(
    ("run", "outcome"),
    (
        (PrismaManager.build_request_log_indexes, False),
        (PrismaManager.start_request_log_index_build, None),
    ),
    ids=("wait-for-the-build", "start-the-build"),
)
def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising(monkeypatch, run, outcome):
    """The migration job exits non-zero and a serving proxy keeps booting when the extras
    package that owns the index build is not installed."""
    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None)

    assert run() is outcome


class _SlowEngine:
    """Engine double whose requests only answer once the test lets them, recording what the engine saw."""

    def __init__(self) -> None:
        self.gates: dict[str, asyncio.Event] = {}  # mutable-ok: one gate per request the test lets through
        self.queries_finished = 0
        self.requests_cancelled = 0
        self.rolled_back: list[str] = []  # mutable-ok: test records
        self.process = object()

    def gate(self, request: str) -> asyncio.Event:
        return self.gates.setdefault(request, asyncio.Event())

    async def _answer_once_released(self, request: str) -> None:
        try:
            await self.gate(request).wait()
        except asyncio.CancelledError:
            self.requests_cancelled += 1
            raise

    async def query(self, content: str, *, tx_id: str | None) -> object:
        await self._answer_once_released(content)
        self.queries_finished += 1
        return {"content": content, "tx_id": tx_id}

    async def start_transaction(self, *, content: str) -> str:
        await self._answer_once_released("start")
        return "tx-opened-after-caller-left"

    async def commit_transaction(self, tx_id: str) -> None:
        raise AssertionError(f"unexpected commit of {tx_id}")

    async def rollback_transaction(self, tx_id: str) -> None:
        self.rolled_back.append(tx_id)


def _tracked(engine: _SlowEngine) -> tuple[_TrackedPrismaEngine, _PrismaDrainTracker]:
    tracker: Final = _PrismaDrainTracker()
    return _TrackedPrismaEngine(cast(Any, engine), tracker), tracker


@pytest.mark.asyncio
async def test_cancelling_a_query_caller_lets_the_engine_finish_the_statement() -> None:
    engine: Final = _SlowEngine()
    tracked, tracker = _tracked(engine)
    caller: Final = asyncio.create_task(tracked.query("x", tx_id=None))
    await asyncio.sleep(0)
    caller.cancel()
    with pytest.raises(asyncio.CancelledError):
        await caller
    drained: Final = asyncio.create_task(tracker.wait_until_drained())
    await asyncio.sleep(0)
    assert engine.requests_cancelled == 0
    assert not drained.done()
    engine.gate("x").set()
    await asyncio.wait_for(drained, timeout=1)
    assert engine.queries_finished == 1


@pytest.mark.asyncio
async def test_transaction_opened_after_its_caller_was_cancelled_is_rolled_back() -> None:
    engine: Final = _SlowEngine()
    tracked, tracker = _tracked(engine)
    caller: Final = asyncio.create_task(tracked.start_transaction(content="{}"))
    await asyncio.sleep(0)
    caller.cancel()
    with pytest.raises(asyncio.CancelledError):
        await caller
    engine.gate("start").set()
    await asyncio.wait_for(tracker.wait_until_drained(), timeout=1)
    assert engine.requests_cancelled == 0
    assert engine.rolled_back == ["tx-opened-after-caller-left"]


@pytest.mark.asyncio
async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain_open() -> None:
    engine: Final = _SlowEngine()
    tracked, tracker = _tracked(engine)
    first: Final = asyncio.create_task(tracked.query("first", tx_id=None))
    await asyncio.sleep(0)
    await asyncio.sleep(0)
    drained: Final = asyncio.create_task(tracker.wait_until_drained())
    await asyncio.sleep(0)
    second: Final = asyncio.create_task(tracked.query("second", tx_id=None))
    engine.gate("first").set()
    await first
    await asyncio.sleep(0)
    assert not drained.done()
    engine.gate("second").set()
    await second
    await asyncio.wait_for(drained, timeout=1)
    assert engine.queries_finished == 2