mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(proxy): support PostgreSQL 10 request log index discovery
Use the operator-class vector to count index keys across PostgreSQL versions and cover missing, equivalent and INCLUDE indexes with live catalog tests Type the backported tests, construct environments without mutation and inject startup dependencies while preserving automatic index builds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo
parent
d30a921dfe
commit
7085864085
8 changed files with 350 additions and 289 deletions
|
|
@ -80,7 +80,7 @@ _EQUIVALENT_INDEXES_SQL: Final = (
|
|||
"SELECT i.relname AS name, x.indisvalid AS valid "
|
||||
"FROM pg_index x JOIN pg_class i ON i.oid = x.indexrelid JOIN pg_am am ON am.oid = i.relam "
|
||||
"WHERE x.indrelid = to_regclass(%s) AND i.relname <> %s AND am.amname = 'btree' AND NOT x.indisunique "
|
||||
"AND x.indexprs IS NULL AND x.indpred IS NULL AND x.indnkeyatts = x.indnatts "
|
||||
"AND x.indexprs IS NULL AND x.indpred IS NULL AND cardinality(x.indclass::oid[]) = x.indnatts "
|
||||
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indoption::int2[]) o WHERE o <> 0) "
|
||||
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indclass::oid[]) c JOIN pg_opclass oc ON oc.oid = c WHERE NOT oc.opcdefault) "
|
||||
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indcollation::oid[]) WITH ORDINALITY c(coll, ord) "
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import re
|
|||
import subprocess
|
||||
import sys
|
||||
import urllib.parse as urlparse
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
|
|
@ -702,6 +702,21 @@ class ProxyInitializationHelpers:
|
|||
return multiproc_dir
|
||||
|
||||
|
||||
def _setup_database_with_request_log_indexes(
|
||||
skip_server_startup: bool,
|
||||
*,
|
||||
setup: Callable[[], bool],
|
||||
build: Callable[[], bool],
|
||||
start: Callable[[], None],
|
||||
) -> bool:
|
||||
if not setup():
|
||||
return False
|
||||
if skip_server_startup:
|
||||
return build()
|
||||
start()
|
||||
return True
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument("cli_args", nargs=-1)
|
||||
@click.option("--host", default="0.0.0.0", help="Host for the server to listen on.", envvar="HOST")
|
||||
|
|
@ -1414,15 +1429,15 @@ def run_server(
|
|||
"LiteLLM versions contend for the same DB.\033[0m"
|
||||
)
|
||||
try:
|
||||
migrated: Final = PrismaManager.setup_database(
|
||||
use_migrate=not use_prisma_db_push,
|
||||
use_v2_resolver=use_v2_resolver,
|
||||
setup_ok: Final = _setup_database_with_request_log_indexes(
|
||||
skip_server_startup,
|
||||
setup=lambda: PrismaManager.setup_database(
|
||||
use_migrate=not use_prisma_db_push,
|
||||
use_v2_resolver=use_v2_resolver,
|
||||
),
|
||||
build=PrismaManager.build_request_log_indexes,
|
||||
start=PrismaManager.start_request_log_index_build,
|
||||
)
|
||||
setup_ok: Final = migrated and (
|
||||
not skip_server_startup or PrismaManager.build_request_log_indexes()
|
||||
)
|
||||
if migrated and not skip_server_startup:
|
||||
PrismaManager.start_request_log_index_build()
|
||||
except RuntimeError as e:
|
||||
# Raised on unrecoverable migration errors: the v2
|
||||
# resolver's non-idempotent failures and permission
|
||||
|
|
|
|||
|
|
@ -825,6 +825,45 @@ def _create_index(database_url: str, name: str, table: str, columns: str) -> int
|
|||
return _index_oids(database_url)[name]
|
||||
|
||||
|
||||
@requires_db
|
||||
@pytest.mark.parametrize(
|
||||
("definition", "matches"),
|
||||
(
|
||||
(None, False),
|
||||
('("litellm_call_id", "api_key")', True),
|
||||
('("litellm_call_id") INCLUDE ("api_key")', False),
|
||||
),
|
||||
ids=("missing", "equivalent", "included-column"),
|
||||
)
|
||||
def test_equivalent_index_lookup_works_without_postgres_11_catalog_columns(
|
||||
definition: str | None, matches: bool
|
||||
) -> None:
|
||||
schema: Final = f"index_catalog_{uuid.uuid4().hex}"
|
||||
target: Final = RequestLogIndex("LiteLLM_SpendLogs", "managed_lookup", '("litellm_call_id", "api_key")')
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
if definition is not None and "INCLUDE" in definition and conn.info.server_version < 110000:
|
||||
pytest.skip("INCLUDE indexes require PostgreSQL 11")
|
||||
conn.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
|
||||
try:
|
||||
conn.execute(
|
||||
sql.SQL('CREATE TABLE {}."LiteLLM_SpendLogs" (litellm_call_id text, api_key text)').format(
|
||||
sql.Identifier(schema)
|
||||
)
|
||||
)
|
||||
if definition is not None:
|
||||
conn.execute(
|
||||
sql.SQL('CREATE INDEX operator_lookup ON {}."LiteLLM_SpendLogs" {}').format(
|
||||
sql.Identifier(schema), sql.SQL(definition)
|
||||
)
|
||||
)
|
||||
found: Final = request_log_indexes._equivalent_indexes(conn, schema, target.table, target.name, target)
|
||||
assert tuple((index.name, index.valid) for index in found) == (
|
||||
(("operator_lookup", True),) if matches else ()
|
||||
)
|
||||
finally:
|
||||
conn.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_valid_index_of_the_same_definition_under_another_name_is_renamed_instead_of_rebuilt(
|
||||
scratch_database: str,
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import os
|
|||
import signal
|
||||
import sys
|
||||
import urllib.parse
|
||||
from collections.abc import Callable
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -464,7 +465,9 @@ def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashi
|
|||
),
|
||||
ids=("wait-for-the-build", "start-the-build"),
|
||||
)
|
||||
def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising(monkeypatch, run, outcome):
|
||||
def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising(
|
||||
monkeypatch: pytest.MonkeyPatch, run: Callable[[], bool | None], outcome: bool | None
|
||||
) -> None:
|
||||
"""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)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@ import inspect
|
|||
import os
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import click
|
||||
|
|
@ -22,6 +24,7 @@ from uvicorn.config import LOOP_FACTORIES
|
|||
from uvicorn.importer import import_from_string
|
||||
|
||||
from litellm.proxy import proxy_cli
|
||||
from litellm.proxy.db import prisma_client
|
||||
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
|
||||
|
||||
|
||||
|
|
@ -2126,21 +2129,29 @@ class TestQueryEngineReaperWiring:
|
|||
class TestRunServerDbSetup:
|
||||
"""Tests for run_server's prisma setup_database behavior."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_setup_database(self, monkeypatch: pytest.MonkeyPatch) -> MagicMock:
|
||||
setup: Final = MagicMock()
|
||||
manager: Final = SimpleNamespace(
|
||||
setup_database=setup,
|
||||
build_request_log_indexes=lambda: True,
|
||||
start_request_log_index_build=lambda: None,
|
||||
)
|
||||
monkeypatch.setattr(prisma_client, "PrismaManager", manager)
|
||||
return setup
|
||||
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_use_prisma_db_push_flag_behavior(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
):
|
||||
mock_should_update_schema: MagicMock,
|
||||
mock_check_schema_diff: MagicMock,
|
||||
mock_atexit_register: MagicMock,
|
||||
mock_subprocess_run: MagicMock,
|
||||
mock_setup_database: MagicMock,
|
||||
) -> None:
|
||||
"""Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter"""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
|
|
@ -2150,19 +2161,17 @@ class TestRunServerDbSetup:
|
|||
# Mock should_update_prisma_schema to return True (so setup_database gets called)
|
||||
mock_should_update_schema.return_value = True
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
mock_proxy_module: Final = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL")
|
||||
clean_env: Final = {
|
||||
**{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")},
|
||||
"DATABASE_URL": "postgresql://test:test@localhost:5432/test",
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
|
|
@ -2206,39 +2215,38 @@ class TestRunServerDbSetup:
|
|||
mock_setup_database.assert_called_with(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch(
|
||||
"litellm.proxy.db.check_migration.check_prisma_schema_diff"
|
||||
) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch(
|
||||
"litellm.proxy.db.prisma_client.should_update_prisma_schema"
|
||||
) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
def test_migrations_run_when_the_prisma_cli_is_not_on_path(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
tmp_path,
|
||||
capsys,
|
||||
):
|
||||
mock_should_update_schema: MagicMock,
|
||||
mock_check_schema_diff: MagicMock,
|
||||
mock_atexit_register: MagicMock,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
mock_setup_database: MagicMock,
|
||||
) -> None:
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
mock_should_update_schema.return_value = True
|
||||
empty_bin = tmp_path / "emptybin"
|
||||
empty_bin: Final = tmp_path / "emptybin"
|
||||
empty_bin.mkdir()
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
mock_proxy_module: Final = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL")
|
||||
clean_env: Final = {
|
||||
**{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")},
|
||||
"DATABASE_URL": "postgresql://test:test@localhost:5432/test",
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
clean_env["PATH"] = str(empty_bin)
|
||||
|
||||
with (
|
||||
|
|
@ -2279,23 +2287,20 @@ class TestRunServerDbSetup:
|
|||
)
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_startup_fails_when_db_setup_fails(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
arguments,
|
||||
env,
|
||||
exits,
|
||||
capsys,
|
||||
):
|
||||
mock_should_update_schema: MagicMock,
|
||||
mock_check_schema_diff: MagicMock,
|
||||
mock_atexit_register: MagicMock,
|
||||
mock_subprocess_run: MagicMock,
|
||||
arguments: tuple[str, ...],
|
||||
env: dict[str, str],
|
||||
exits: bool,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
mock_setup_database: MagicMock,
|
||||
) -> None:
|
||||
"""A failed PrismaManager.setup_database exits 1 by default; only an explicit
|
||||
--no-enforce_prisma_migration_check or ENFORCE_PRISMA_MIGRATION_CHECK=false keeps the proxy booting"""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
|
@ -2304,20 +2309,22 @@ class TestRunServerDbSetup:
|
|||
mock_should_update_schema.return_value = True
|
||||
mock_setup_database.return_value = False
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
mock_proxy_module: Final = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL", "ENFORCE_PRISMA_MIGRATION_CHECK")
|
||||
clean_env: Final = {
|
||||
**{
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL", "ENFORCE_PRISMA_MIGRATION_CHECK")
|
||||
},
|
||||
"DATABASE_URL": "postgresql://test:test@localhost:5432/test",
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
outcome = pytest.raises(SystemExit) if exits else nullcontext()
|
||||
outcome: Final = pytest.raises(SystemExit) if exits else nullcontext()
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {**clean_env, **env}, clear=True),
|
||||
|
|
@ -2346,19 +2353,16 @@ class TestRunServerDbSetup:
|
|||
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_startup_exits_on_non_postgres_database_url(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
):
|
||||
mock_should_update_schema: MagicMock,
|
||||
mock_check_schema_diff: MagicMock,
|
||||
mock_atexit_register: MagicMock,
|
||||
mock_subprocess_run: MagicMock,
|
||||
mock_setup_database: MagicMock,
|
||||
) -> None:
|
||||
"""A sqlite DATABASE_URL must exit immediately, before any prisma call,
|
||||
instead of stalling on a migration against the postgresql-only schema."""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
|
@ -2366,19 +2370,17 @@ class TestRunServerDbSetup:
|
|||
mock_subprocess_run.return_value = MagicMock(returncode=0)
|
||||
mock_should_update_schema.return_value = True
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
mock_proxy_module: Final = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL")
|
||||
clean_env: Final = {
|
||||
**{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")},
|
||||
"DATABASE_URL": "sqlite:///data/litellm.db",
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "sqlite:///data/litellm.db"
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
|
|
@ -2391,28 +2393,23 @@ class TestRunServerDbSetup:
|
|||
),
|
||||
):
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
run_server.main(
|
||||
["--local", "--skip_server_startup"], standalone_mode=False
|
||||
)
|
||||
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
|
||||
assert exc_info.value.code == 1
|
||||
mock_setup_database.assert_not_called()
|
||||
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_v2_migration_resolver_opts_in_via_env_var(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
capsys,
|
||||
):
|
||||
mock_should_update_schema: MagicMock,
|
||||
mock_check_schema_diff: MagicMock,
|
||||
mock_atexit_register: MagicMock,
|
||||
mock_subprocess_run: MagicMock,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
mock_setup_database: MagicMock,
|
||||
) -> None:
|
||||
"""USE_V2_MIGRATION_RESOLVER=true must select the v2 resolver.
|
||||
|
||||
The Helm migrations Job runs `python litellm/proxy/prisma_migration.py`,
|
||||
|
|
@ -2425,19 +2422,17 @@ class TestRunServerDbSetup:
|
|||
mock_should_update_schema.return_value = True
|
||||
mock_setup_database.return_value = True
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
mock_proxy_module: Final = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL")
|
||||
clean_env: Final = {
|
||||
**{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")},
|
||||
"DATABASE_URL": "postgresql://test:test@localhost:5432/test",
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
clean_env["USE_V2_MIGRATION_RESOLVER"] = "true"
|
||||
|
||||
with (
|
||||
|
|
@ -2450,9 +2445,7 @@ class TestRunServerDbSetup:
|
|||
},
|
||||
),
|
||||
):
|
||||
run_server.main(
|
||||
["--local", "--skip_server_startup"], standalone_mode=False
|
||||
)
|
||||
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
|
||||
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
|
||||
assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out
|
||||
|
|
@ -2477,12 +2470,7 @@ class TestRunServerDbSetup:
|
|||
def test_resolve_v2_migration_resolver(self, use_legacy_flag, env_value, expected):
|
||||
from litellm.proxy.proxy_cli import resolve_v2_migration_resolver
|
||||
|
||||
assert (
|
||||
resolve_v2_migration_resolver(
|
||||
use_legacy_flag=use_legacy_flag, env_value=env_value
|
||||
)
|
||||
is expected
|
||||
)
|
||||
assert resolve_v2_migration_resolver(use_legacy_flag=use_legacy_flag, env_value=env_value) is expected
|
||||
|
||||
def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self):
|
||||
from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli
|
||||
|
|
@ -2491,19 +2479,16 @@ class TestRunServerDbSetup:
|
|||
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_legacy_resolver_flag_reaches_database_setup(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
):
|
||||
mock_should_update_schema: MagicMock,
|
||||
mock_check_schema_diff: MagicMock,
|
||||
mock_atexit_register: MagicMock,
|
||||
mock_subprocess_run: MagicMock,
|
||||
mock_setup_database: MagicMock,
|
||||
) -> None:
|
||||
"""--use_legacy_migration_resolver must reach the database setup call.
|
||||
|
||||
The resolver decision itself is covered mock-free above; this is the
|
||||
|
|
@ -2515,19 +2500,21 @@ class TestRunServerDbSetup:
|
|||
mock_should_update_schema.return_value = True
|
||||
mock_setup_database.return_value = True
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
mock_proxy_module: Final = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER")
|
||||
clean_env: Final = {
|
||||
**{
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER")
|
||||
},
|
||||
"DATABASE_URL": "postgresql://test:test@localhost:5432/test",
|
||||
}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
|
|
@ -2551,73 +2538,51 @@ class TestRunServerDbSetup:
|
|||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("arguments", "migrated", "exits", "waits_for_the_build"),
|
||||
("skip_server_startup", "migrated", "build_succeeds", "calls_expected", "succeeds"),
|
||||
(
|
||||
(("--local", "--skip_server_startup"), True, True, True),
|
||||
(("--local",), True, False, False),
|
||||
(("--local",), False, True, False),
|
||||
(True, True, False, ("setup", "build"), False),
|
||||
(True, True, True, ("setup", "build"), True),
|
||||
(False, True, False, ("setup", "start"), True),
|
||||
(False, False, False, ("setup",), False),
|
||||
(True, False, False, ("setup",), False),
|
||||
),
|
||||
ids=(
|
||||
"failed-build",
|
||||
"migration-job",
|
||||
"serving-proxy",
|
||||
"serving-proxy-whose-migrations-failed",
|
||||
"failed-migration-job",
|
||||
),
|
||||
ids=("migration-job", "serving-proxy", "serving-proxy-whose-migrations-failed"),
|
||||
)
|
||||
@patch("uvicorn.run")
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True)
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes", return_value=False)
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.start_request_log_index_build")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=True)
|
||||
def test_the_migration_job_waits_for_the_index_build_and_a_serving_proxy_starts_it_in_the_background(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_start_build,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
mock_uvicorn_run,
|
||||
arguments,
|
||||
migrated,
|
||||
exits,
|
||||
waits_for_the_build,
|
||||
):
|
||||
"""`--skip_server_startup` is the migration job: it waits for the index build after the
|
||||
migrations and exits 1 when one could not be built. A serving proxy that ran the
|
||||
migrations starts the build in the background and serves whatever the build does; one
|
||||
whose migrations failed exits 1 under `--enforce_prisma_migration_check` and starts no build."""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
skip_server_startup: bool,
|
||||
migrated: bool,
|
||||
build_succeeds: bool,
|
||||
calls_expected: tuple[str, ...],
|
||||
succeeds: bool,
|
||||
) -> None:
|
||||
calls: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
|
||||
mock_setup_database.return_value = migrated
|
||||
mock_subprocess_run.return_value = MagicMock(returncode=0)
|
||||
mock_proxy_module = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
def setup() -> bool:
|
||||
calls.put("setup")
|
||||
return migrated
|
||||
|
||||
def build() -> bool:
|
||||
calls.put("build")
|
||||
return build_succeeds
|
||||
|
||||
def start() -> None:
|
||||
calls.put("start")
|
||||
|
||||
assert (
|
||||
proxy_cli._setup_database_with_request_log_indexes(
|
||||
skip_server_startup, setup=setup, build=build, start=start
|
||||
)
|
||||
is succeeds
|
||||
)
|
||||
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
outcome = pytest.raises(SystemExit) if exits else nullcontext()
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module},
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
|
||||
) as mock_get_args,
|
||||
outcome as exc_info,
|
||||
):
|
||||
mock_get_args.return_value = {"app": "litellm.proxy.proxy_server:app", "host": "localhost", "port": 8000}
|
||||
run_server.main([*arguments, "--enforce_prisma_migration_check"], standalone_mode=False)
|
||||
|
||||
assert (exc_info is not None and exc_info.value.code == 1) is exits
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
|
||||
assert mock_build_indexes.call_count == int(migrated and waits_for_the_build)
|
||||
assert mock_start_build.call_count == int(migrated and not waits_for_the_build)
|
||||
assert tuple(calls.get_nowait() for _ in calls_expected) == calls_expected
|
||||
assert calls.empty()
|
||||
|
||||
|
||||
# --- Module-level helpers for worker startup hook tests ---
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@ import glob
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -1250,7 +1252,6 @@ class TestJWTKeyMappingCascade:
|
|||
)
|
||||
|
||||
|
||||
|
||||
class TestStripPrismaQueryParams:
|
||||
"""The psycopg URL the job connects with is derived from the Prisma-dialect
|
||||
DATABASE_URL, whose TLS params mean something else to libpq."""
|
||||
|
|
@ -1261,47 +1262,53 @@ class TestStripPrismaQueryParams:
|
|||
|
||||
return dict(parse_qsl(urlparse(url).query))
|
||||
|
||||
def test_prisma_ca_sslcert_becomes_sslrootcert_with_verify_full(self):
|
||||
url = "postgresql://u:p@writer:5432/db?schema=public&sslmode=require&sslcert=/tmp/pinned.pem&sslaccept=strict"
|
||||
def test_prisma_ca_sslcert_becomes_sslrootcert_with_verify_full(self) -> None:
|
||||
url: Final = (
|
||||
"postgresql://u:p@writer:5432/db?schema=public&sslmode=require&sslcert=/tmp/pinned.pem&sslaccept=strict"
|
||||
)
|
||||
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
|
||||
assert self._query(cleaned) == {"sslmode": "verify-full", "sslrootcert": "/tmp/pinned.pem"}
|
||||
assert cleaned.startswith("postgresql://u:p@writer:5432/db?")
|
||||
|
||||
@pytest.mark.parametrize("sslmode", ["prefer", "require"])
|
||||
@pytest.mark.parametrize("sslaccept", ["strict", "unknown-mode-prisma-treats-as-strict"])
|
||||
def test_strict_verifies_chain_and_hostname_whatever_sslmode_prisma_was_given(self, sslmode, sslaccept):
|
||||
url = f"postgresql://writer/db?sslmode={sslmode}&sslcert=/certs/ca.pem&sslaccept={sslaccept}"
|
||||
def test_strict_verifies_chain_and_hostname_whatever_sslmode_prisma_was_given(
|
||||
self, sslmode: str, sslaccept: str
|
||||
) -> None:
|
||||
url: Final = f"postgresql://writer/db?sslmode={sslmode}&sslcert=/certs/ca.pem&sslaccept={sslaccept}"
|
||||
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
|
||||
assert self._query(cleaned) == {"sslmode": "verify-full", "sslrootcert": "/certs/ca.pem"}
|
||||
|
||||
def test_strict_with_tls_disabled_stays_off(self):
|
||||
url = "postgresql://writer/db?sslmode=disable&sslcert=/certs/ca.pem&sslaccept=strict"
|
||||
def test_strict_with_tls_disabled_stays_off(self) -> None:
|
||||
url: Final = "postgresql://writer/db?sslmode=disable&sslcert=/certs/ca.pem&sslaccept=strict"
|
||||
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
|
||||
assert self._query(cleaned) == {"sslmode": "disable"}
|
||||
|
||||
@pytest.mark.parametrize("sslaccept", ["&sslaccept=accept_invalid_certs", ""])
|
||||
def test_without_strict_the_ca_is_dropped_so_libpq_checks_nothing_like_prisma(self, sslaccept):
|
||||
url = f"postgresql://writer/db?sslmode=require&sslcert=/certs/ca.pem{sslaccept}"
|
||||
def test_without_strict_the_ca_is_dropped_so_libpq_checks_nothing_like_prisma(self, sslaccept: str) -> None:
|
||||
url: Final = f"postgresql://writer/db?sslmode=require&sslcert=/certs/ca.pem{sslaccept}"
|
||||
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
|
||||
assert self._query(cleaned) == {"sslmode": "require"}
|
||||
|
||||
def test_a_ca_alone_without_strict_or_sslmode_leaves_libpq_its_defaults(self):
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params("postgresql://writer/db?sslcert=/certs/ca.pem")
|
||||
def test_a_ca_alone_without_strict_or_sslmode_leaves_libpq_its_defaults(self) -> None:
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params("postgresql://writer/db?sslcert=/certs/ca.pem")
|
||||
|
||||
assert cleaned == "postgresql://writer/db"
|
||||
|
||||
def test_a_libpq_client_certificate_pair_is_left_alone(self):
|
||||
url = "postgresql://writer/db?sslmode=verify-full&sslrootcert=/ca.pem&sslcert=/client.crt&sslkey=/client.key"
|
||||
def test_a_libpq_client_certificate_pair_is_left_alone(self) -> None:
|
||||
url: Final = (
|
||||
"postgresql://writer/db?sslmode=verify-full&sslrootcert=/ca.pem&sslcert=/client.crt&sslkey=/client.key"
|
||||
)
|
||||
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
|
||||
assert self._query(cleaned) == {
|
||||
"sslmode": "verify-full",
|
||||
|
|
@ -1310,17 +1317,17 @@ class TestStripPrismaQueryParams:
|
|||
"sslkey": "/client.key",
|
||||
}
|
||||
|
||||
def test_an_explicit_sslrootcert_wins_over_the_prisma_sslcert(self):
|
||||
url = "postgresql://writer/db?sslmode=require&sslrootcert=/ca.pem&sslcert=/pinned.pem&sslaccept=strict"
|
||||
def test_an_explicit_sslrootcert_wins_over_the_prisma_sslcert(self) -> None:
|
||||
url: Final = "postgresql://writer/db?sslmode=require&sslrootcert=/ca.pem&sslcert=/pinned.pem&sslaccept=strict"
|
||||
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
|
||||
assert self._query(cleaned) == {"sslmode": "verify-full", "sslrootcert": "/ca.pem"}
|
||||
|
||||
def test_prisma_only_params_are_dropped_and_plain_urls_pass_through(self):
|
||||
url = "postgresql://u:p@pooler:6543/db?schema=tenant&pgbouncer=true&connection_limit=5&connect_timeout=3"
|
||||
def test_prisma_only_params_are_dropped_and_plain_urls_pass_through(self) -> None:
|
||||
url: Final = "postgresql://u:p@pooler:6543/db?schema=tenant&pgbouncer=true&connection_limit=5&connect_timeout=3"
|
||||
|
||||
cleaned = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
cleaned: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
|
||||
assert cleaned == "postgresql://u:p@pooler:6543/db?connect_timeout=3"
|
||||
assert (
|
||||
|
|
@ -1334,66 +1341,77 @@ class TestBuildRequestLogIndexes:
|
|||
the migrations target, waits for it, and reports its result."""
|
||||
|
||||
@pytest.fixture
|
||||
def builds(self):
|
||||
return []
|
||||
def builds(self) -> SimpleQueue[tuple[str, str]]:
|
||||
return SimpleQueue()
|
||||
|
||||
@pytest.fixture
|
||||
def build(self, builds):
|
||||
def build(self, builds: SimpleQueue[tuple[str, str]]) -> Callable[[str, str], bool]:
|
||||
def record(database_url: str, schema: str) -> bool:
|
||||
builds.append((database_url, schema))
|
||||
builds.put((database_url, schema))
|
||||
return True
|
||||
|
||||
return record
|
||||
|
||||
def test_the_build_gets_the_direct_url_without_prisma_params_and_the_prisma_schema(self, monkeypatch, builds, build):
|
||||
def test_the_build_gets_the_direct_url_without_prisma_params_and_the_prisma_schema(
|
||||
self, monkeypatch: pytest.MonkeyPatch, builds: SimpleQueue[tuple[str, str]], build: Callable[[str, str], bool]
|
||||
) -> None:
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@pooler:6543/db?schema=tenant&pgbouncer=true")
|
||||
monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@primary:5432/db?connection_limit=1")
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
|
||||
|
||||
assert builds == [("postgresql://u:p@primary:5432/db", "tenant")]
|
||||
assert builds.get_nowait() == ("postgresql://u:p@primary:5432/db", "tenant")
|
||||
assert builds.empty()
|
||||
|
||||
def test_the_build_defaults_to_the_database_url_and_the_public_schema(self, monkeypatch, builds, build):
|
||||
def test_the_build_defaults_to_the_database_url_and_the_public_schema(
|
||||
self, monkeypatch: pytest.MonkeyPatch, builds: SimpleQueue[tuple[str, str]], build: Callable[[str, str], bool]
|
||||
) -> None:
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
|
||||
|
||||
assert builds == [("postgresql://u:p@primary:5432/db", "public")]
|
||||
assert builds.get_nowait() == ("postgresql://u:p@primary:5432/db", "public")
|
||||
assert builds.empty()
|
||||
|
||||
def test_a_build_that_leaves_indexes_missing_is_reported_so_the_job_reruns(self, monkeypatch):
|
||||
def test_a_build_that_leaves_indexes_missing_is_reported_so_the_job_reruns(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=lambda url, schema: False) is False
|
||||
|
||||
def test_without_a_database_url_nothing_is_built(self, monkeypatch, builds, build):
|
||||
def test_without_a_database_url_nothing_is_built(
|
||||
self, monkeypatch: pytest.MonkeyPatch, builds: SimpleQueue[tuple[str, str]], build: Callable[[str, str], bool]
|
||||
) -> None:
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
|
||||
|
||||
assert builds == []
|
||||
assert builds.empty()
|
||||
|
||||
|
||||
class TestStartRequestLogIndexBuild:
|
||||
"""A serving proxy that ran the migrations starts the index build on a daemon thread
|
||||
and goes on to serve while it runs."""
|
||||
|
||||
def test_the_build_runs_on_a_daemon_thread_that_does_not_hold_up_the_caller(self):
|
||||
def test_the_build_runs_on_a_daemon_thread_that_does_not_hold_up_the_caller(self) -> None:
|
||||
release: Final = threading.Event()
|
||||
builds: Final[list[str]] = [] # mutable-ok: the builder thread hands back the thread it ran on
|
||||
builds: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
|
||||
def build() -> bool:
|
||||
assert release.wait(5), "the caller never came back from start_request_log_index_build"
|
||||
builds.append(threading.current_thread().name)
|
||||
builds.put(threading.current_thread().name)
|
||||
return True
|
||||
|
||||
thread: Final = ProxyExtrasDBManager.start_request_log_index_build(build=build)
|
||||
|
||||
assert builds == [], "the build ran before start_request_log_index_build returned"
|
||||
assert builds.empty(), "the build ran before start_request_log_index_build returned"
|
||||
assert thread.daemon is True
|
||||
release.set()
|
||||
thread.join(5)
|
||||
assert builds == ["litellm-request-log-indexes"]
|
||||
assert builds.get_nowait() == "litellm-request-log-indexes"
|
||||
assert builds.empty()
|
||||
|
||||
|
||||
class TestRunMigrationJob:
|
||||
|
|
@ -1401,14 +1419,16 @@ class TestRunMigrationJob:
|
|||
result deciding whether the job reports success."""
|
||||
|
||||
@pytest.fixture
|
||||
def calls(self):
|
||||
return []
|
||||
def calls(self) -> SimpleQueue[tuple[str, bool, bool] | tuple[str]]:
|
||||
return SimpleQueue()
|
||||
|
||||
@pytest.fixture
|
||||
def setup(self, calls):
|
||||
def record(result: bool):
|
||||
def setup(
|
||||
self, calls: SimpleQueue[tuple[str, bool, bool] | tuple[str]]
|
||||
) -> Callable[[bool], Callable[[bool, bool], bool]]:
|
||||
def record(result: bool) -> Callable[[bool, bool], bool]:
|
||||
def setup_database(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
calls.append(("setup", use_migrate, use_v2_resolver))
|
||||
calls.put(("setup", use_migrate, use_v2_resolver))
|
||||
return result
|
||||
|
||||
return setup_database
|
||||
|
|
@ -1416,50 +1436,68 @@ class TestRunMigrationJob:
|
|||
return record
|
||||
|
||||
@pytest.fixture
|
||||
def build(self, calls):
|
||||
def record(result: bool):
|
||||
def build(self, calls: SimpleQueue[tuple[str, bool, bool] | tuple[str]]) -> Callable[[bool], Callable[[], bool]]:
|
||||
def record(result: bool) -> Callable[[], bool]:
|
||||
def build_request_log_indexes() -> bool:
|
||||
calls.append(("build",))
|
||||
calls.put(("build",))
|
||||
return result
|
||||
|
||||
return build_request_log_indexes
|
||||
|
||||
return record
|
||||
|
||||
def test_the_job_builds_the_indexes_after_the_migrations_succeed(self, calls, setup, build):
|
||||
def test_the_job_builds_the_indexes_after_the_migrations_succeed(
|
||||
self,
|
||||
calls: SimpleQueue[tuple[str, bool, bool] | tuple[str]],
|
||||
setup: Callable[[bool], Callable[[bool, bool], bool]],
|
||||
build: Callable[[bool], Callable[[], bool]],
|
||||
) -> None:
|
||||
assert ProxyExtrasDBManager.run_migration_job(True, False, setup=setup(True), build=build(True)) is True
|
||||
|
||||
assert calls == [("setup", True, False), ("build",)]
|
||||
assert (calls.get_nowait(), calls.get_nowait()) == (("setup", True, False), ("build",))
|
||||
assert calls.empty()
|
||||
|
||||
def test_the_job_fails_without_building_when_the_migrations_fail(self, calls, setup, build):
|
||||
def test_the_job_fails_without_building_when_the_migrations_fail(
|
||||
self,
|
||||
calls: SimpleQueue[tuple[str, bool, bool] | tuple[str]],
|
||||
setup: Callable[[bool], Callable[[bool, bool], bool]],
|
||||
build: Callable[[bool], Callable[[], bool]],
|
||||
) -> None:
|
||||
assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(False), build=build(True)) is False
|
||||
|
||||
assert calls == [("setup", True, True)]
|
||||
assert calls.get_nowait() == ("setup", True, True)
|
||||
assert calls.empty()
|
||||
|
||||
def test_the_job_fails_when_an_index_could_not_be_built(self, calls, setup, build):
|
||||
def test_the_job_fails_when_an_index_could_not_be_built(
|
||||
self,
|
||||
calls: SimpleQueue[tuple[str, bool, bool] | tuple[str]],
|
||||
setup: Callable[[bool], Callable[[bool, bool], bool]],
|
||||
build: Callable[[bool], Callable[[], bool]],
|
||||
) -> None:
|
||||
assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(True), build=build(False)) is False
|
||||
|
||||
assert calls == [("setup", True, True), ("build",)]
|
||||
assert (calls.get_nowait(), calls.get_nowait()) == (("setup", True, True), ("build",))
|
||||
assert calls.empty()
|
||||
|
||||
|
||||
class TestMigrationJobOwnedDrift:
|
||||
JOB_INDEXES = (
|
||||
JOB_INDEXES: Final = (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
|
||||
"\n-- CreateIndex\n"
|
||||
'CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
|
||||
)
|
||||
|
||||
def test_a_plain_spend_logs_table_only_loses_the_migration_job_indexes(self):
|
||||
filtered = ProxyExtrasDBManager._filter_migration_job_owned_drift(
|
||||
def test_a_plain_spend_logs_table_only_loses_the_migration_job_indexes(self) -> None:
|
||||
filtered: Final = ProxyExtrasDBManager._filter_migration_job_owned_drift(
|
||||
_PARTITIONED_DRIFT_SQL + self.JOB_INDEXES, partitioned=False
|
||||
)
|
||||
assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
|
||||
assert "LiteLLM_SpendLogs_api_key_startTime_idx" not in filtered
|
||||
assert 'PRIMARY KEY ("request_id")' in filtered
|
||||
|
||||
def test_a_partitioned_spend_logs_table_also_loses_its_partitioning_artifacts(self):
|
||||
filtered = ProxyExtrasDBManager._filter_migration_job_owned_drift(
|
||||
def test_a_partitioned_spend_logs_table_also_loses_its_partitioning_artifacts(self) -> None:
|
||||
filtered: Final = ProxyExtrasDBManager._filter_migration_job_owned_drift(
|
||||
_PARTITIONED_DRIFT_SQL + self.JOB_INDEXES, partitioned=True
|
||||
)
|
||||
assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
|
||||
|
|
|
|||
|
|
@ -34,13 +34,13 @@ def _prisma_indexes_of(schema: str, model: str) -> frozenset[str]:
|
|||
|
||||
|
||||
class TestTheIndexList:
|
||||
def test_every_migration_job_index_is_declared_in_the_prisma_schema_under_the_same_name(self):
|
||||
def test_every_migration_job_index_is_declared_in_the_prisma_schema_under_the_same_name(self) -> None:
|
||||
schema: Final = SCHEMA.read_text()
|
||||
for index in REQUEST_LOG_INDEXES:
|
||||
assert index.name in _prisma_indexes_of(schema, index.table), index
|
||||
|
||||
@pytest.mark.parametrize("name", INERT_MIGRATIONS)
|
||||
def test_the_migrations_that_used_to_build_these_indexes_run_no_sql(self, name: str):
|
||||
def test_the_migrations_that_used_to_build_these_indexes_run_no_sql(self, name: str) -> None:
|
||||
assert is_inert_migration((PACKAGE / "migrations" / name / "migration.sql").read_text())
|
||||
|
||||
|
||||
|
|
@ -66,7 +66,7 @@ class TestIsInertMigration:
|
|||
"lowercase",
|
||||
),
|
||||
)
|
||||
def test_comments_and_a_select_1_alone_are_inert(self, script: str):
|
||||
def test_comments_and_a_select_1_alone_are_inert(self, script: str) -> None:
|
||||
assert is_inert_migration(script) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -80,40 +80,40 @@ class TestIsInertMigration:
|
|||
),
|
||||
ids=("select-2", "select-from", "index-after-comment", "alter-after-block-comment", "drop-after-select-1"),
|
||||
)
|
||||
def test_any_statement_is_not_inert(self, script: str):
|
||||
def test_any_statement_is_not_inert(self, script: str) -> None:
|
||||
assert is_inert_migration(script) is False
|
||||
|
||||
|
||||
class TestPartitionIndexName:
|
||||
def test_a_partition_gets_the_name_postgres_would_give_an_inherited_index(self):
|
||||
def test_a_partition_gets_the_name_postgres_would_give_an_inherited_index(self) -> None:
|
||||
assert CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p2026_09") == (
|
||||
"LiteLLM_SpendLogs_p2026_09_litellm_call_id_idx"
|
||||
)
|
||||
|
||||
def test_an_index_not_prefixed_by_its_table_keeps_its_whole_name(self):
|
||||
index = RequestLogIndex("LiteLLM_SpendLogs", "call_id_lookup", '("litellm_call_id")')
|
||||
def test_an_index_not_prefixed_by_its_table_keeps_its_whole_name(self) -> None:
|
||||
index: Final = RequestLogIndex("LiteLLM_SpendLogs", "call_id_lookup", '("litellm_call_id")')
|
||||
assert index.partition_index_name("LiteLLM_SpendLogs_pdefault") == "LiteLLM_SpendLogs_pdefault_call_id_lookup"
|
||||
|
||||
def test_a_long_name_is_cut_to_63_bytes_with_a_digest_that_keeps_partitions_apart(self):
|
||||
first = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_09")
|
||||
second = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_10")
|
||||
def test_a_long_name_is_cut_to_63_bytes_with_a_digest_that_keeps_partitions_apart(self) -> None:
|
||||
first: Final = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_09")
|
||||
second: Final = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_10")
|
||||
assert len(first.encode()) == 63 and len(second.encode()) == 63
|
||||
assert first != second
|
||||
assert first.startswith("LiteLLM_SpendLogs_p") and first[-9] == "_"
|
||||
|
||||
def test_the_byte_limit_counts_multibyte_characters(self):
|
||||
name = CALL_ID_INDEX.partition_index_name("é" * 40)
|
||||
def test_the_byte_limit_counts_multibyte_characters(self) -> None:
|
||||
name: Final = CALL_ID_INDEX.partition_index_name("é" * 40)
|
||||
assert len(name.encode()) <= 63 and len(name) < 63
|
||||
|
||||
|
||||
class TestColumns:
|
||||
def test_the_columns_are_the_quoted_names_of_the_definition_in_order(self):
|
||||
index = RequestLogIndex(
|
||||
def test_the_columns_are_the_quoted_names_of_the_definition_in_order(self) -> None:
|
||||
index: Final = RequestLogIndex(
|
||||
"LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'
|
||||
)
|
||||
assert index.columns == ("api_key", "startTime")
|
||||
|
||||
def test_every_migration_job_index_names_at_least_one_column(self):
|
||||
def test_every_migration_job_index_names_at_least_one_column(self) -> None:
|
||||
assert all(index.columns for index in REQUEST_LOG_INDEXES)
|
||||
|
||||
|
||||
|
|
@ -127,23 +127,23 @@ DRIFT_WITH_BOTH_INDEXES: Final = (
|
|||
|
||||
|
||||
class TestFilterRequestLogIndexDiff:
|
||||
def test_a_drift_script_that_only_creates_the_migration_job_indexes_becomes_empty(self):
|
||||
def test_a_drift_script_that_only_creates_the_migration_job_indexes_becomes_empty(self) -> None:
|
||||
assert filter_request_log_index_diff(DRIFT_WITH_BOTH_INDEXES) == ""
|
||||
|
||||
def test_other_statements_survive_with_the_migration_job_indexes_removed(self):
|
||||
def test_other_statements_survive_with_the_migration_job_indexes_removed(self) -> None:
|
||||
other: Final = '-- AlterTable\nALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;\n'
|
||||
filtered = filter_request_log_index_diff(other + DRIFT_WITH_BOTH_INDEXES)
|
||||
filtered: Final = filter_request_log_index_diff(other + DRIFT_WITH_BOTH_INDEXES)
|
||||
assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;' in filtered
|
||||
assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
|
||||
assert "LiteLLM_SpendLogs_api_key_startTime_idx" not in filtered
|
||||
|
||||
def test_an_index_of_another_name_on_spend_logs_is_kept(self):
|
||||
def test_an_index_of_another_name_on_spend_logs_is_kept(self) -> None:
|
||||
sql: Final = 'CREATE INDEX "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user");\n'
|
||||
assert filter_request_log_index_diff(sql) == sql
|
||||
|
||||
def test_a_drop_of_a_migration_job_index_is_kept_for_the_operator_to_see(self):
|
||||
def test_a_drop_of_a_migration_job_index_is_kept_for_the_operator_to_see(self) -> None:
|
||||
sql: Final = 'DROP INDEX "LiteLLM_SpendLogs_litellm_call_id_idx";\n'
|
||||
assert filter_request_log_index_diff(sql) == sql
|
||||
|
||||
def test_an_empty_script_stays_empty(self):
|
||||
def test_an_empty_script_stays_empty(self) -> None:
|
||||
assert filter_request_log_index_diff("") == ""
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for conditional DD
|
|||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -212,43 +213,43 @@ class TestIndexesOnLogTables:
|
|||
the whole build and a concurrent one fails on a partitioned parent, so the migration job
|
||||
(litellm_proxy_extras/request_log_indexes.py) builds those instead."""
|
||||
|
||||
def test_the_original_spend_log_index_statement_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
def test_the_original_spend_log_index_statement_is_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("api_key", "startTime");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_the_original_concurrent_call_id_index_statement_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
def test_the_original_concurrent_call_id_index_statement_is_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = (
|
||||
'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("litellm_call_id");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_unique_index_with_if_not_exists_on_error_logs_is_flagged(self, tmp_path):
|
||||
sql = 'CREATE UNIQUE INDEX IF NOT EXISTS "ix" ON "LiteLLM_ErrorLogs" ("request_id");'
|
||||
def test_unique_index_with_if_not_exists_on_error_logs_is_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = 'CREATE UNIQUE INDEX IF NOT EXISTS "ix" ON "LiteLLM_ErrorLogs" ("request_id");'
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_ErrorLogs"',)
|
||||
|
||||
def test_a_unique_concurrent_index_on_error_logs_is_flagged(self, tmp_path):
|
||||
sql = 'CREATE UNIQUE INDEX CONCURRENTLY "ix" ON "LiteLLM_ErrorLogs" ("request_id");'
|
||||
def test_a_unique_concurrent_index_on_error_logs_is_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = 'CREATE UNIQUE INDEX CONCURRENTLY "ix" ON "LiteLLM_ErrorLogs" ("request_id");'
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_ErrorLogs"',)
|
||||
|
||||
def test_lowercase_schema_qualified_and_only_forms_are_flagged(self, tmp_path):
|
||||
sql = (
|
||||
def test_lowercase_schema_qualified_and_only_forms_are_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = (
|
||||
'create index on "public"."LiteLLM_SpendLogs" ("api_key");\n'
|
||||
'CREATE INDEX "ix" ON ONLY "LiteLLM_SpendLogs" ("api_key");\n'
|
||||
'CREATE INDEX CONCURRENTLY "iy" ON "public"."LiteLLM_SpendLogs" ("api_key");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) * 3
|
||||
|
||||
def test_a_comment_between_on_and_the_table_is_flagged(self, tmp_path):
|
||||
sql = 'CREATE INDEX CONCURRENTLY "ix" ON /* table */ "LiteLLM_SpendLogs" ("api_key");'
|
||||
def test_a_comment_between_on_and_the_table_is_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = 'CREATE INDEX CONCURRENTLY "ix" ON /* table */ "LiteLLM_SpendLogs" ("api_key");'
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_a_concurrent_index_with_comments_and_line_breaks_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
def test_a_concurrent_index_with_comments_and_line_breaks_is_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX CONCURRENTLY IF NOT EXISTS "ix"\n'
|
||||
' ON "LiteLLM_SpendLogs" /* partitioned in some deployments */\n'
|
||||
|
|
@ -256,34 +257,34 @@ class TestIndexesOnLogTables:
|
|||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_indexes_on_a_non_log_table_pass_concurrent_or_not(self, tmp_path):
|
||||
sql = (
|
||||
def test_indexes_on_a_non_log_table_pass_concurrent_or_not(self, tmp_path: Path) -> None:
|
||||
sql: Final = (
|
||||
'CREATE INDEX "ix" ON "LiteLLM_VerificationToken" ("token");\n'
|
||||
'CREATE INDEX CONCURRENTLY "iy" ON "LiteLLM_VerificationToken" ("token");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_an_index_run_by_execute_is_flagged(self, tmp_path):
|
||||
sql = 'DO $$ BEGIN EXECUTE \'CREATE INDEX "ix" ON "LiteLLM_SpendLogs" ("api_key")\'; END $$;'
|
||||
def test_an_index_run_by_execute_is_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = 'DO $$ BEGIN EXECUTE \'CREATE INDEX "ix" ON "LiteLLM_SpendLogs" ("api_key")\'; END $$;'
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_a_marker_does_not_exempt_the_index(self, tmp_path):
|
||||
sql = (
|
||||
def test_a_marker_does_not_exempt_the_index(self, tmp_path: Path) -> None:
|
||||
sql: Final = (
|
||||
'-- data-migration-ok: table is empty at this point\nCREATE INDEX "ix" ON "LiteLLM_SpendLogs" ("api_key");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_a_marker_on_a_rewrite_still_leaves_the_index_below_it_flagged(self, tmp_path):
|
||||
sql = (
|
||||
def test_a_marker_on_a_rewrite_still_leaves_the_index_below_it_flagged(self, tmp_path: Path) -> None:
|
||||
sql: Final = (
|
||||
"-- data-migration-ok: one row\n"
|
||||
'UPDATE "LiteLLM_SpendLogs" SET "api_key" = \'k\' WHERE "request_id" = \'r\';\n'
|
||||
'CREATE INDEX CONCURRENTLY "ix" ON "LiteLLM_SpendLogs" ("api_key");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_render_points_at_the_migration_job_index_list(self, tmp_path):
|
||||
sql = 'CREATE INDEX CONCURRENTLY "ix" ON "public"."LiteLLM_SpendLogs" ("api_key");'
|
||||
rendered = _scan(tmp_path, sql)[0].render()
|
||||
def test_render_points_at_the_migration_job_index_list(self, tmp_path: Path) -> None:
|
||||
sql: Final = 'CREATE INDEX CONCURRENTLY "ix" ON "public"."LiteLLM_SpendLogs" ("api_key");'
|
||||
rendered: Final = _scan(tmp_path, sql)[0].render()
|
||||
assert "20260101000000_fixture/migration.sql:1" in rendered
|
||||
assert "blocks writes until the build finishes, or fails on a partitioned table" in rendered
|
||||
assert "REQUEST_LOG_INDEXES in litellm_proxy_extras/request_log_indexes.py" in rendered
|
||||
|
|
@ -292,7 +293,7 @@ class TestIndexesOnLogTables:
|
|||
"name",
|
||||
("20260823000000_add_spend_logs_api_key_starttime_index", "20260831120001_spend_logs_litellm_call_id_index"),
|
||||
)
|
||||
def test_the_inert_index_migrations_scan_clean_without_a_grandfather(self, name):
|
||||
def test_the_inert_index_migrations_scan_clean_without_a_grandfather(self, name: str) -> None:
|
||||
assert checker.scan_migration(checker.MIGRATIONS_DIR / name) == ()
|
||||
assert name not in checker.GRANDFATHERED
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue