litellm/tests/integration/database/test_reader_writer_regeneration.py

import json
import os
import uuid
from concurrent.futures import ThreadPoolExecutor
from hashlib import sha256
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit

import httpx
import psycopg
import pytest
from integration._support.client import Gateway, delete_key_if_present, eventually, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from psycopg import sql
from pydantic import BaseModel, JsonValue, TypeAdapter


@pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants")
def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None:
    role: Final = f"integration_reader_{uuid.uuid4().hex}"
    url: Final = os.environ["DATABASE_URL"]
    parsed: Final = urlsplit(url)
    reader_url: Final = urlunsplit(
        parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}")
    )
    with psycopg.connect(url, autocommit=True) as admin:
        admin.execute(
            sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format(
                sql.Identifier(role)
            )
        )
        try:
            admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role)))
            admin.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(sql.Identifier(role)))
            admin.execute(sql.SQL("ALTER ROLE {} SET default_transaction_read_only = on").format(sql.Identifier(role)))
            with psycopg.connect(reader_url, autocommit=True) as reader:
                assert reader.execute("SHOW transaction_read_only").fetchone() == ("on",)
                with pytest.raises(psycopg.errors.ReadOnlySqlTransaction):
                    reader.execute('UPDATE "LiteLLM_VerificationToken" SET blocked = true WHERE false')
            with owned_proxy(gateway, tmp_path, {"DATABASE_URL_READ_REPLICA": reader_url}) as candidate:
                assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), (
                    "Candidate reader was never connected"
                )
                with gateway.scenario() as scenario:
                    model: Final = scenario.model()
                    outside: Final = scenario.model()
                    old: Final = string_value(candidate.post("/key/generate", {"models": [outside]})["key"])
                    new: Final = f"sk-integration-{uuid.uuid4().hex}"
                    scenario.cleanups.callback(delete_key_if_present, gateway, old)
                    scenario.cleanups.callback(delete_key_if_present, gateway, new)
                    old_hash: Final = sha256(old.encode()).hexdigest()
                    before: Final = candidate.request(
                        "POST",
                        "/v1/chat/completions",
                        {"model": model, "messages": [{"role": "user", "content": "no grant yet"}]},
                        key=old,
                    )
                    assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", (
                        before.text
                    )
                    response: Final = candidate.request(
                        "POST",
                        "/v1/access_group",
                        {
                            "access_group_name": f"integration-{uuid.uuid4().hex}",
                            "access_model_names": [model],
                            "assigned_key_ids": [old_hash],
                        },
                    )
                    assert response.status_code == 201, response.text
                    group: Final = string_value(response.json()["access_group_id"])
                    try:
                        with psycopg.connect(url) as blocker, ThreadPoolExecutor(max_workers=1) as executor:
                            blocker.execute('LOCK TABLE "LiteLLM_AccessGroupTable" IN ACCESS EXCLUSIVE MODE')
                            pending: Final = executor.submit(candidate.request, "GET", f"/v1/access_group/{group}")
                            try:
                                reached: Final = eventually(
                                    lambda: read_rows(
                                        "SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) "
                                        "AND usename=%s AND query LIKE 'SELECT%%'",
                                        (blocker.info.backend_pid, role),
                                    ),
                                    bool,
                                    seconds=3,
                                )
                                assert reached == [{"usename": role}]
                            finally:
                                blocker.rollback()
                            selected: Final = pending.result(timeout=5)
                            assert selected.status_code == 200 and selected.json()["access_group_id"] == group, (
                                selected.text
                            )
                        assert candidate.chat(model, key=old)["usage"]["total_tokens"] == 40
                        regenerated: Final = candidate.post(
                            "/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"}
                        )
                        assert regenerated["key"] == new
                        new_hash: Final = sha256(new.encode()).hexdigest()
                        assert new != old
                        assert read_rows(
                            'SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)
                        ) == [{"assigned_key_ids": [new_hash]}]
                        assert read_rows(
                            'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)',
                            ([old_hash, new_hash],),
                        ) == [{"token": new_hash, "access_group_ids": [group]}]
                        assert candidate.chat(model, key=new)["usage"]["total_tokens"] == 40
                        assert candidate.chat(outside, key=new)["usage"]["total_tokens"] == 40
                        denied: Final = candidate.request(
                            "POST",
                            "/v1/chat/completions",
                            {"model": model, "messages": [{"role": "user", "content": "rotated key"}]},
                            key=old,
                        )
                        assert (
                            denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db"
                        ), denied.text
                    finally:
                        deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group}")
                        assert deleted.status_code == 204, deleted.text
                        assert (
                            read_rows(
                                'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s',
                                (group,),
                            )
                            == []
                        )
        finally:
            admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role)))
            admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role)))
    assert read_rows("SELECT rolname FROM pg_roles WHERE rolname=%s", (role,)) == []


class _Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class _AssistantMessage(BaseModel):
    content: str


class _Choice(BaseModel):
    message: _AssistantMessage


class _ChatResponse(BaseModel):
    id: str
    choices: tuple[_Choice, ...]
    usage: _Usage


class _GenerateKeyResponse(BaseModel):
    key: str


class _Error(BaseModel):
    type: str


class _ErrorResponse(BaseModel):
    error: _Error


class _SpendLog(BaseModel):
    api_key: str
    spend: float


class _KeySpend(BaseModel):
    spend: float


class _DeprecatedKey(BaseModel):
    active_token_id: str
    grace_seconds: float


_CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue])
_EXPECTED_CHAT: Final = ("POST", "/v1/chat/completions")


def _chat_body(model: str, prompt: str) -> dict[str, JsonValue]:
    return {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
    }


def _wire_reply(identity: str) -> Reply:
    return Reply(
        body=json.dumps(
            {
                "id": identity,
                "object": "chat.completion",
                "created": 1,
                "model": "integration-keys-auth-grace",
                "choices": [
                    {
                        "index": 0,
                        "message": {"role": "assistant", "content": "wire response"},
                        "finish_reason": "stop",
                    }
                ],
                "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
            }
        ).encode()
    )


def _chat(gateway: Gateway, model: str, key: str, prompt: str) -> httpx.Response:
    return gateway.request("POST", "/v1/chat/completions", _chat_body(model, prompt), key=key)


def _assert_success(response: httpx.Response) -> _ChatResponse:
    assert response.status_code == 200, response.text
    payload: Final = _ChatResponse.model_validate_json(response.content)
    assert payload.choices[0].message.content == "wire response", response.text
    assert payload.usage == _Usage(prompt_tokens=20, completion_tokens=20, total_tokens=40), response.text
    return payload


@pytest.mark.parametrize("spelling", ("body", "path", "plaintext_path"))
def test_regenerate_with_a_grace_period_keeps_the_old_key_serving_billed_to_the_new_row(
    gateway: Gateway, spelling: str
) -> None:
    assert os.environ.get("LITELLM_KEY_ROTATION_GRACE_PERIOD") is None
    backend: Final = "integration-keys-auth-grace"
    prompts: Final = (
        "grace old key " + uuid.uuid4().hex,
        "grace new key " + uuid.uuid4().hex,
        "revoked second key " + uuid.uuid4().hex,
        "final generated key " + uuid.uuid4().hex,
    )

    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/v1/chat/completions"
        assert request.headers["authorization"] == "Bearer synthetic-openai-key"
        body: Final = _CHAT_BODY.validate_json(request.body)
        assert body in tuple(_chat_body(backend, prompt) for prompt in prompts), body
        return _wire_reply("chatcmpl-" + uuid.uuid4().hex)

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(
            model=f"openai/{backend}",
            api_base=f"{wire.url}/v1",
            api_key="synthetic-openai-key",
            input_cost_per_token=0.001,
            output_cost_per_token=0.002,
        )
        created: Final = gateway.request("POST", "/key/generate", {"models": [model]})
        assert created.status_code == 200, created.text
        key_one: Final = _GenerateKeyResponse.model_validate_json(created.content).key
        scenario.cleanups.callback(delete_key_if_present, gateway, key_one)
        hash_one: Final = sha256(key_one.encode()).hexdigest()

        path_key_one: Final = hash_one if spelling == "path" else key_one
        generated_two: Final = (
            gateway.request("POST", "/key/regenerate", {"key": key_one, "grace_period": "1h"})
            if spelling == "body"
            else gateway.request("POST", f"/key/{path_key_one}/regenerate", {"grace_period": "1h"})
        )
        assert generated_two.status_code == 200, generated_two.text
        key_two: Final = _GenerateKeyResponse.model_validate_json(generated_two.content).key
        scenario.cleanups.callback(delete_key_if_present, gateway, key_two)
        hash_two: Final = sha256(key_two.encode()).hexdigest()
        old_response: Final = _assert_success(_chat(gateway, model, key_one, prompts[0]))
        new_response: Final = _assert_success(_chat(gateway, model, key_two, prompts[1]))
        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2

        old_spend: Final = TypeAdapter(tuple[_SpendLog, ...]).validate_python(
            eventually(
                lambda: read_rows(
                    'SELECT api_key, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
                    (old_response.id,),
                ),
                lambda rows: len(rows) == 1,
                seconds=70,
            )
        )
        new_spend: Final = TypeAdapter(tuple[_SpendLog, ...]).validate_python(
            eventually(
                lambda: read_rows(
                    'SELECT api_key, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
                    (new_response.id,),
                ),
                lambda rows: len(rows) == 1,
                seconds=70,
            )
        )
        assert old_spend[0].api_key == hash_two
        assert old_spend[0].spend > 0
        assert new_spend[0].api_key == hash_two
        assert new_spend[0].spend > 0
        key_spend: Final = TypeAdapter(tuple[_KeySpend, ...]).validate_python(
            eventually(
                lambda: read_rows(
                    'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s',
                    (hash_two,),
                ),
                lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0,
                seconds=70,
            )
        )
        assert key_spend[0].spend > 0
        deprecated: Final = TypeAdapter(tuple[_DeprecatedKey, ...]).validate_python(
            read_rows(
                "SELECT active_token_id, EXTRACT(EPOCH FROM revoke_at - now())::float8 AS grace_seconds "
                'FROM "LiteLLM_DeprecatedVerificationToken" WHERE token=%s',
                (hash_one,),
            )
        )
        assert len(deprecated) == 1
        assert deprecated[0].active_token_id == hash_two
        assert deprecated[0].grace_seconds == pytest.approx(3600, abs=120)

        path_key_two: Final = hash_two if spelling == "path" else key_two
        generated_three: Final = (
            gateway.request("POST", "/key/regenerate", {"key": key_two})
            if spelling == "body"
            else gateway.request("POST", f"/key/{path_key_two}/regenerate", {})
        )
        assert generated_three.status_code == 200, generated_three.text
        key_three: Final = _GenerateKeyResponse.model_validate_json(generated_three.content).key
        scenario.cleanups.callback(delete_key_if_present, gateway, key_three)
        refused: Final = _chat(gateway, model, key_two, prompts[2])
        assert refused.status_code == 401, refused.text
        assert _ErrorResponse.model_validate_json(refused.content).error.type == "token_not_found_in_db", refused.text
        assert wire.drain() == (), "the revoked second key reached the upstream"
        assert (
            _assert_success(_chat(gateway, model, key_three, prompts[3])).choices[0].message.content == "wire response"
        )
        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,)