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

import json
import os
import signal
import sys
import urllib.parse
from unittest.mock import AsyncMock, MagicMock, Mock, patch

import pytest
from fastapi.testclient import TestClient



from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, 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"]


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]
    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]


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]
    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, bypassing the operator's enforce_prisma_migration_check choice.
    """
    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)

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