litellm/tests/e2e/perf/setup/ui_query_fixture.py

from __future__ import annotations

import hashlib
import os
import time
from datetime import datetime, timedelta
from typing import Final

import psycopg
from pydantic import AwareDatetime, BaseModel, ConfigDict, TypeAdapter

from setup.ui_query_isolation import validate_target
from setup.upgrade_worker import read_state, release_indexes


class UiFixture(BaseModel):
    model_config = ConfigDict(frozen=True)

    anchor: AwareDatetime
    api_key: str


SEEDED_TEAMS: Final = 100
SEEDED_USERS: Final = 2000
SEEDED_KEYS: Final = 5000
SEEDED_MODELS: Final = 200
DAILY_TABLES: Final = ("LiteLLM_DailyUserSpend", "LiteLLM_DailyTeamSpend")
_SEED_SQL: Final = (
    """
    INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, members_with_roles, metadata, model_spend, model_max_budget)
    SELECT 'perf-team-' || n, 'perf-team-' || n, '[]', '{}', '{}', '{}' FROM generate_series(0, %(teams)s - 1) AS n
    ON CONFLICT (team_id) DO NOTHING
    """,
    """
    INSERT INTO "LiteLLM_UserTable" (user_id, user_alias, user_email, user_role, team_id, metadata, model_spend, model_max_budget)
    SELECT 'perf-user-' || n, 'perf-user-' || n, 'perf-user-' || n || '@example.com', 'internal_user',
        'perf-team-' || (n %% %(teams)s), '{}', '{}', '{}'
    FROM generate_series(0, %(users)s - 1) AS n
    ON CONFLICT (user_id) DO NOTHING
    """,
    """
    INSERT INTO "LiteLLM_VerificationToken" (token, key_name, key_alias, team_id, user_id, aliases, config, permissions,
        metadata, model_spend, model_max_budget, created_at, updated_at)
    SELECT md5('perf-key:' || n) || md5('perf-key-tail:' || n), 'sk-...perf' || n, 'perf-key-' || n,
        'perf-team-' || (n %% %(teams)s), 'perf-user-' || (n %% %(users)s), '{}', '{}', '{}', '{}', '{}', '{}',
        %(anchor)s - make_interval(days => n %% 365), %(anchor)s
    FROM generate_series(0, %(keys)s - 1) AS n
    ON CONFLICT (token) DO NOTHING
    """,
    """
    INSERT INTO "LiteLLM_ProxyModelTable" (model_id, model_name, litellm_params, model_info, created_by, updated_by)
    SELECT 'perf-model-' || n, 'perf-chat-' || (n %% 20),
        jsonb_build_object('model', 'openai/perf-chat', 'api_key', 'perf-seed',
            'api_base', 'os.environ/E2E_PERF_MOCK_OPENAI_URL'),
        jsonb_build_object('id', 'perf-model-' || n, 'db_model', true), 'perf-seed', 'perf-seed'
    FROM generate_series(0, %(models)s - 1) AS n
    ON CONFLICT (model_id) DO NOTHING
    """,
)


def prepare_fixture() -> UiFixture:
    validate_target(os.environ["DATABASE_URL"])
    _finish_upgrade()
    with psycopg.connect(os.environ["DATABASE_URL"], connect_timeout=10) as connection:
        _ = connection.execute("SET statement_timeout='30s'")
        row: Final = connection.execute("SELECT spec->>'anchor', rows FROM litellm_perf_seed.checkpoint").fetchone()
        assert row is not None and TypeAdapter(int).validate_python(row[1]) > 0, "Missing completed seed rows"
        seed_ui_tables(connection, anchor=TypeAdapter(datetime).validate_python(row[0]))
        fixture: Final = UiFixture(
            anchor=TypeAdapter(datetime).validate_python(row[0]),
            api_key=hashlib.md5(b"perf-key:0", usedforsecurity=False).hexdigest()
            + hashlib.md5(b"perf-key-tail:0", usedforsecurity=False).hexdigest(),
        )
        return fixture


def seed_ui_tables(connection: psycopg.Connection, anchor: datetime) -> None:
    """Give the Keys, Teams, Users and Models pages realistic row counts on the disposable restore.

    Names follow the SpendLogs seed (perf-key-N aliases, perf-team-N, perf-user-N) so the Logs filters and
    Team Usage line up with real spend. The Daily* rollups come from the snapshot itself and are only checked.
    """
    window: Final = {
        "since": (anchor - timedelta(days=7)).strftime("%Y-%m-%d"),
        "until": anchor.strftime("%Y-%m-%d"),
    }
    for table in DAILY_TABLES:
        row = connection.execute(f'SELECT count(*) FROM "{table}" WHERE date BETWEEN %(since)s AND %(until)s', window)
        assert _count(row.fetchone()) > 0, f"{table} has no rows in the week before the anchor"
    params: Final = {
        "teams": SEEDED_TEAMS,
        "users": SEEDED_USERS,
        "keys": SEEDED_KEYS,
        "models": SEEDED_MODELS,
        "anchor": anchor,
    }
    for statement in _SEED_SQL:
        _ = connection.execute(statement, params)
    connection.commit()


def _count(row: tuple[object, ...] | None) -> int:
    assert row is not None
    return TypeAdapter(int).validate_python(row[0])


def _released_schema(connection: psycopg.Connection) -> bool:
    """A released tag's migrations image upgrades the restore (schema and request-log indexes) before the
    stack starts, so there is no upgrade worker state to wait on."""
    row: Final = connection.execute("""SELECT to_regclass('public."_LiteLLMPerfUpgrade"') IS NULL""").fetchone()
    return row is None or TypeAdapter(bool).validate_python(row[0])


def _finish_upgrade() -> None:
    if os.environ.get("E2E_DATABASE_UPGRADE") != "1":
        return
    with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True, connect_timeout=10) as connection:
        if _released_schema(connection):
            return
        release_indexes(connection)
        deadline: Final = time.monotonic() + 6 * 3600
        while time.monotonic() < deadline:
            match read_state(connection).phase:
                case "failed":
                    raise AssertionError("Snapshot upgrade failed")
                case "complete":
                    return
                case _:
                    pass
            time.sleep(5)
    raise AssertionError("Snapshot indexes did not finish before UI measurements")