litellm/tests/e2e/migrations/test_startup.py

from contextlib import ExitStack
from itertools import accumulate, repeat
from typing import Final

import pytest

from .checks import (
    COMPLETE,
    FATAL,
    GATED,
    assert_completed,
    assert_completed_after_lock_timeouts,
    start_replicas,
)
from .containers import Containers, failed, ready, until, waiting
from .database import GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user
from .startup_models import Migration
from e2e_metadata import Domain, Subject, meta

pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]


def union_waiting_identities(
    observed: frozenset[str], snapshot: frozenset[str]
) -> frozenset[str]:
    return observed | snapshot


class TestMigrationStartup:
    @pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False)))
    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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,),)

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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)

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    def test_waiters_retry_lock_timeouts_until_contention_clears(
        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)
                waiting_identities: Final = accumulate(
                    (database.waiting_identities((PRISMA_LOCK, GATE_KEY)) for _ in repeat(None)),
                    union_waiting_identities,
                    initial=frozenset(),
                )
                follower_identities: Final = frozenset(follower.identity for follower in followers)
                until(
                    "both followers entered migration locking",
                    lambda: follower_identities <= next(waiting_identities),
                )
                until(
                    "a replica logged a migration lock-timeout retry",
                    lambda: any(
                        "timed out waiting for a table lock" in replica.logs()
                        for replica in (owner, *followers)
                    ),
                )
                waiting((owner, *followers), 20)
            ready((owner, *followers), database)
            assert_completed_after_lock_timeouts(database, GATED)

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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)

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    def test_lock_timeout_exhaustion_then_restart(self, containers: Containers, database: Database) -> None:
        with database.lock():
            with containers.start(
                database, (GATED,), environment={"LITELLM_MIGRATION_DDL_LOCK_TIMEOUT": "2"}
            ) as replica:
                until("owner blocked in migration SQL", lambda: bool(database.blocked()))
                failed((replica,), "attempts that made no progress")
            assert not database.exists("migration_effect")
            assert database.query(
                "SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND finished_at IS NOT NULL",
                (GATED.name,),
            ) == ((0,),)
        with containers.start(database, (GATED,)) as restarted:
            ready((restarted,), database)
        assert_completed_after_lock_timeouts(database, GATED)

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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,),)

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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))
    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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'))

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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

    @meta(
        Subject(
            domain=Domain.DB,
        )
    )
    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")