diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 6c31e8364a9..b2548f08645 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -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) " diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 2f94b115120..cf8f2d230e1 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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 diff --git a/tests/proxy_migration_tests/test_request_log_indexes.py b/tests/proxy_migration_tests/test_request_log_indexes.py index 23e4adce477..9c34b60f044 100644 --- a/tests/proxy_migration_tests/test_request_log_indexes.py +++ b/tests/proxy_migration_tests/test_request_log_indexes.py @@ -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, diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index a34d3c0c27d..f753ddce3c1 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -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) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 046f841945c..0847c8534cb 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -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 --- diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 219d6b6afbd..dc1b67669a2 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -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 diff --git a/tests/unit/litellm_proxy_extras/test_request_log_indexes.py b/tests/unit/litellm_proxy_extras/test_request_log_indexes.py index 5cf7593c5cd..4db1e407fe6 100644 --- a/tests/unit/litellm_proxy_extras/test_request_log_indexes.py +++ b/tests/unit/litellm_proxy_extras/test_request_log_indexes.py @@ -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("") == "" diff --git a/tests/unit/test_check_migrations_no_data_rewrites.py b/tests/unit/test_check_migrations_no_data_rewrites.py index fb1053ff887..a2e9bb6887d 100644 --- a/tests/unit/test_check_migrations_no_data_rewrites.py +++ b/tests/unit/test_check_migrations_no_data_rewrites.py @@ -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