litellm/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
yuneng-jiang ae0e8a20db
fix(ci): run the migration DDL guard, and stop it reading comments as SQL (#37791)
* fix(ci): make the migration DDL guard run, and stop it reading comments as SQL

TestMigrationSQLIdempotency requires guarded DDL across litellm-proxy-extras
and has never run in any job, so the convention eroded quietly. Four of its
assertions fail today, and it was allowlisted rather than wired up because
fixing the migrations is not an option: Prisma checksums an applied migration,
so editing one breaks `migrate deploy` for every existing install.

Two things were wrong with the guard itself. It scanned raw lines, so Prisma's
own `-- CREATE INDEX CONCURRENTLY ...` explanations counted as the statements
they describe, which is two of the reported migrations. And it had no way to
say "these predate the rule", so the only options were editing immutable files
or leaving the whole file unrun.

Comments are now stripped before matching, on the drop-column rule too, and the
migrations that already violate are named once in _PRE_GUARD_MIGRATIONS. The
rules bind everything after them, so a new migration with bare CREATE TABLE,
ADD COLUMN, CREATE INDEX or an unguarded ADD CONSTRAINT now fails a check
instead of landing unnoticed.

That set is 14 migrations, not the 13 previously recorded, measured after
comment-stripping. It can only shrink: a test fails if an entry names no
migration on disk, and another fails if an entry no longer violates anything.

The file now runs as a proxy-extras shard and comes off the coverage allowlist.

* fix(ci): strip block comments in the migration guard too

Prisma opens a destructive migration with a /* Warnings: You are about to
drop the column ... */ header. Nothing in the tree trips a rule on that text
today, but it is prose about a statement rather than the statement, and the
line-comment fix left the class open. Bodies are blanked rather than removed
so the reported line number still points at the real statement.
2026-08-22 22:57:16 -07:00

477 lines
20 KiB
Python

import glob
import os
import re
import sys
import pytest
sys.path.insert(
0,
os.path.abspath(
os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras")
),
)
from litellm_proxy_extras.utils import ProxyExtrasDBManager
# Path to the migrations directory
_MIGRATIONS_DIR = os.path.abspath(
os.path.join(
os.path.dirname(__file__),
"../../litellm-proxy-extras/litellm_proxy_extras/migrations",
)
)
def test_custom_prisma_dir(monkeypatch):
import tempfile
# create a temp directory
temp_dir = tempfile.mkdtemp()
monkeypatch.setenv("LITELLM_MIGRATION_DIR", temp_dir)
## Check if the prisma dir is the temp directory
assert ProxyExtrasDBManager._get_prisma_dir() == temp_dir
## Check if the schema.prisma file is in the temp directory
schema_path = os.path.join(temp_dir, "schema.prisma")
assert os.path.exists(schema_path)
## Check if the migrations dir is in the temp directory
migrations_dir = os.path.join(temp_dir, "migrations")
assert os.path.exists(migrations_dir)
class TestPermissionErrorDetection:
"""Test cases for permission error detection in Prisma migrations"""
def test_is_permission_error_postgres_42501(self):
"""Test detection of PostgreSQL 42501 error code (insufficient privilege)"""
error_message = "Database error code: 42501 - permission denied for table users"
assert ProxyExtrasDBManager._is_permission_error(error_message) is True
def test_is_permission_error_must_be_owner(self):
"""Test detection of 'must be owner of table' error"""
error_message = "ERROR: must be owner of table my_table"
assert ProxyExtrasDBManager._is_permission_error(error_message) is True
def test_is_permission_error_permission_denied_schema(self):
"""Test detection of 'permission denied for schema' error"""
error_message = "permission denied for schema public"
assert ProxyExtrasDBManager._is_permission_error(error_message) is True
def test_is_permission_error_permission_denied_table(self):
"""Test detection of 'permission denied for table' error"""
error_message = "permission denied for table my_table"
assert ProxyExtrasDBManager._is_permission_error(error_message) is True
def test_is_permission_error_must_be_owner_schema(self):
"""Test detection of 'must be owner of schema' error"""
error_message = "must be owner of schema public"
assert ProxyExtrasDBManager._is_permission_error(error_message) is True
def test_is_permission_error_case_insensitive(self):
"""Test that permission error detection is case insensitive"""
error_message = "PERMISSION DENIED FOR TABLE my_table"
assert ProxyExtrasDBManager._is_permission_error(error_message) is True
def test_is_permission_error_negative(self):
"""Test that non-permission errors are not detected as permission errors"""
error_message = "column 'id' already exists"
assert ProxyExtrasDBManager._is_permission_error(error_message) is False
class TestIdempotentErrorDetection:
"""Test cases for idempotent error detection in Prisma migrations"""
def test_is_idempotent_error_already_exists(self):
"""Test detection of generic 'already exists' error"""
error_message = "object already exists"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
def test_is_idempotent_error_column_already_exists(self):
"""Test detection of 'column already exists' error"""
error_message = "column 'email' already exists"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
def test_is_idempotent_error_duplicate_key(self):
"""Test detection of duplicate key violation error"""
error_message = "duplicate key value violates unique constraint"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
def test_is_idempotent_error_relation_already_exists(self):
"""Test detection of 'relation already exists' error"""
error_message = "relation 'users_pkey' already exists"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
def test_is_idempotent_error_constraint_already_exists(self):
"""Test detection of 'constraint already exists' error"""
error_message = "constraint 'fk_user_id' already exists"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
def test_is_idempotent_error_case_insensitive(self):
"""Test that idempotent error detection is case insensitive"""
error_message = "COLUMN 'ID' ALREADY EXISTS"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
def test_is_idempotent_error_does_not_exist(self):
"""Test detection of 'does not exist' error"""
error_message = "ERROR: index 'idx' does not exist"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
def test_is_idempotent_error_negative(self):
"""Test that non-idempotent errors are not detected as idempotent errors"""
error_message = "Database error code: 42501 - permission denied"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False
class TestErrorClassificationPriority:
"""Test cases to ensure errors are correctly classified"""
def test_permission_error_not_classified_as_idempotent(self):
"""Ensure permission errors are not mistakenly classified as idempotent"""
error_message = "Database error code: 42501 - must be owner of table users"
assert ProxyExtrasDBManager._is_permission_error(error_message) is True
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False
def test_idempotent_error_not_classified_as_permission(self):
"""Ensure idempotent errors are not mistakenly classified as permission errors"""
error_message = "column 'created_at' already exists"
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True
assert ProxyExtrasDBManager._is_permission_error(error_message) is False
def test_unknown_error_classified_as_neither(self):
"""Ensure unknown errors are classified as neither permission nor idempotent"""
error_message = "connection timeout"
assert ProxyExtrasDBManager._is_permission_error(error_message) is False
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False
def _get_all_migrations():
"""Return (migration_name, sql_content) pairs for all migrations."""
migration_files = sorted(
glob.glob(os.path.join(_MIGRATIONS_DIR, "*/migration.sql"))
)
results = []
for path in migration_files:
migration_name = os.path.basename(os.path.dirname(path))
with open(path) as f:
results.append((migration_name, f.read()))
return results
_LINE_COMMENT = re.compile(r"--.*$")
_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.DOTALL)
_PRE_GUARD_MIGRATIONS = frozenset({
"20260331000000_add_prompt_environment_and_created_by",
"20260418000000_add_adaptive_router_tables",
"20260429161855_workflow_runs_tables",
"20260605182307_add_timeout_to_mcp_server_table",
"20260626120000_add_mcp_tool_search_enabled",
"20260629000000_add_max_concurrent_requests_to_mcp_server_table",
"20260710000000_add_dcr_bridge_to_mcp_server_table",
"20260713230852_add_key_type_to_litellm_verification_token",
"20260811172448_add_shadow_eval",
"20260813180408_add_shadow_eval_direction",
"20260814000000_add_proxy_worker_heartbeat",
"20260817143646_add_daily_guardrail_usage_units",
"20260818224500_add_shadow_eval_stopped_by",
"20260819000000_shadow_eval_max_budget",
})
def _blanked_block_comments(sql):
"""`sql` with every `/* ... */` body blanked out, newlines kept so lines still count.
Prisma opens a destructive migration with a `/* Warnings: You are about to drop the
column ... */` header, which is prose about the statement rather than the statement.
"""
return _BLOCK_COMMENT.sub(lambda m: re.sub(r"[^\n]", " ", m.group(0)), sql)
def _statements(sql):
"""(line_number, sql) for each line, with comments removed.
Prisma writes its own explanations as `-- CREATE INDEX CONCURRENTLY ...`, which a
raw-line scan reads as the statement it is describing.
"""
return [
(number, _LINE_COMMENT.sub("", line))
for number, line in enumerate(_blanked_block_comments(sql).splitlines(), 1)
]
def _guarded_migrations(all_migrations):
"""Migrations the DDL rules apply to. Prisma checksums an applied migration, so the
ones that predate these rules cannot be edited without breaking `migrate deploy`
for existing installs; they are named once, and the rules bind everything after.
"""
return [
(name, sql) for name, sql in all_migrations if name not in _PRE_GUARD_MIGRATIONS
]
class TestMigrationSQLIdempotency:
"""Ensure all migration SQL files use idempotent DDL (IF [NOT] EXISTS).
Migrations on pre-existing instances can fail when DDL statements assume
the target object doesn't already exist (or still exists for drops).
These tests enforce that all migrations use safe, re-runnable SQL patterns.
"""
@pytest.fixture(scope="class")
def all_migrations(self):
migrations = _get_all_migrations()
assert len(migrations) > 0, (
f"No migrations found. "
f"Check that _MIGRATIONS_DIR ({_MIGRATIONS_DIR}) is correct."
)
return migrations
def test_create_table_uses_if_not_exists(self, all_migrations):
"""CREATE TABLE statements must use IF NOT EXISTS"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
for line_num, line in _statements(sql):
if re.search(
r"CREATE\s+TABLE\s+", line, re.IGNORECASE
) and not re.search(
r"CREATE\s+TABLE\s+IF\s+NOT\s+EXISTS", line, re.IGNORECASE
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert (
not violations
), "CREATE TABLE without IF NOT EXISTS found in migrations:\n" + "\n".join(
violations
)
def test_add_column_uses_if_not_exists(self, all_migrations):
"""ADD COLUMN statements must use IF NOT EXISTS"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
for line_num, line in _statements(sql):
if re.search(r"ADD\s+COLUMN\s+", line, re.IGNORECASE) and not re.search(
r"ADD\s+COLUMN\s+IF\s+NOT\s+EXISTS", line, re.IGNORECASE
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert not violations, (
"ADD COLUMN without IF NOT EXISTS found in recent migrations:\n"
+ "\n".join(violations)
)
def test_drop_column_uses_if_exists(self, all_migrations):
"""DROP COLUMN statements must use IF EXISTS"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
for line_num, line in _statements(sql):
if re.search(
r"DROP\s+COLUMN\s+", line, re.IGNORECASE
) and not re.search(
r"DROP\s+COLUMN\s+IF\s+EXISTS", line, re.IGNORECASE
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert (
not violations
), "DROP COLUMN without IF EXISTS found in recent migrations:\n" + "\n".join(
violations
)
_DROP_COLUMN_ALLOWLIST = {
"20250918083359_drop_spec_version_column_from_mcp_table",
"20260213170952_access_group_change_to_model_name",
"20260224203854_add_agent_object_permissions_table",
}
def test_no_drop_column_statements(self, all_migrations):
"""Migrations must not drop columns — dropping columns is destructive
and can break running application instances during rolling deploys."""
violations = []
for migration_name, sql in all_migrations:
if migration_name in self._DROP_COLUMN_ALLOWLIST:
continue
for line_num, line in _statements(sql):
if re.search(r"DROP\s+COLUMN", line, re.IGNORECASE):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert (
not violations
), "DROP COLUMN found in migrations (destructive, not allowed):\n" + "\n".join(
violations
)
def test_drop_index_uses_if_exists(self, all_migrations):
"""DROP INDEX statements must use IF EXISTS"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
for line_num, line in _statements(sql):
if re.search(r"DROP\s+INDEX\s+", line, re.IGNORECASE) and not re.search(
r"DROP\s+INDEX\s+IF\s+EXISTS", line, re.IGNORECASE
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert (
not violations
), "DROP INDEX without IF EXISTS found in recent migrations:\n" + "\n".join(
violations
)
def test_create_index_uses_if_not_exists(self, all_migrations):
"""CREATE INDEX statements must use IF NOT EXISTS"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
for line_num, line in _statements(sql):
if re.search(
r"CREATE\s+(?:UNIQUE\s+)?INDEX\s+", line, re.IGNORECASE
) and not re.search(
r"CREATE\s+(?:UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?IF\s+NOT\s+EXISTS",
line,
re.IGNORECASE,
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert not violations, (
"CREATE INDEX without IF NOT EXISTS found in recent migrations:\n"
+ "\n".join(violations)
)
def test_rename_column_is_guarded(self, all_migrations):
"""RENAME COLUMN must be inside a DO $$ IF EXISTS block"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
in_do_block = False
for line_num, line in _statements(sql):
if re.search(r"DO\s+\$\$", line, re.IGNORECASE):
in_do_block = True
if re.search(r"END\s+\$\$", line, re.IGNORECASE):
in_do_block = False
if (
re.search(r"RENAME\s+COLUMN\s+", line, re.IGNORECASE)
and not in_do_block
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert not violations, (
"RENAME COLUMN without DO $$ IF EXISTS guard found in migrations:\n"
+ "\n".join(violations)
)
def test_add_constraint_is_guarded(self, all_migrations):
"""ADD CONSTRAINT must be inside a DO $$ IF NOT EXISTS block"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
in_do_block = False
for line_num, line in _statements(sql):
if re.search(r"DO\s+\$\$", line, re.IGNORECASE):
in_do_block = True
if re.search(r"END\s+\$\$", line, re.IGNORECASE):
in_do_block = False
if (
re.search(r"ADD\s+CONSTRAINT\s+", line, re.IGNORECASE)
and not in_do_block
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert not violations, (
"ADD CONSTRAINT without DO $$ IF NOT EXISTS guard found in migrations:\n"
+ "\n".join(violations)
)
def test_drop_constraint_is_guarded(self, all_migrations):
"""DROP CONSTRAINT must be inside a DO $$ IF EXISTS block"""
violations = []
for migration_name, sql in _guarded_migrations(all_migrations):
in_do_block = False
for line_num, line in _statements(sql):
if re.search(r"DO\s+\$\$", line, re.IGNORECASE):
in_do_block = True
if re.search(r"END\s+\$\$", line, re.IGNORECASE):
in_do_block = False
if (
re.search(r"DROP\s+CONSTRAINT\s+", line, re.IGNORECASE)
and not in_do_block
):
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
assert not violations, (
"DROP CONSTRAINT without DO $$ IF EXISTS guard found in migrations:\n"
+ "\n".join(violations)
)
class TestMigrationGuardScope:
"""The guard must ignore SQL comments, exempt only the named pre-guard migrations,
and still fail on a new migration that uses bare DDL."""
_NEW = "20990101000000_a_new_migration"
def _run_rules(self, migrations):
suite = TestMigrationSQLIdempotency()
failures = []
for name in (
"test_create_table_uses_if_not_exists",
"test_add_column_uses_if_not_exists",
"test_create_index_uses_if_not_exists",
"test_add_constraint_is_guarded",
):
try:
getattr(suite, name)(migrations)
except AssertionError:
failures.append(name)
return failures
def test_a_comment_describing_ddl_is_not_the_ddl(self):
sql = '-- CREATE TABLE "Foo" (id TEXT);\n-- ADD COLUMN "bar" TEXT;\n'
assert self._run_rules([(self._NEW, sql)]) == []
def test_a_prisma_warning_block_is_not_the_ddl_it_describes(self):
sql = (
"/*\n"
" Warnings:\n"
"\n"
" - You are about to CREATE TABLE \"Foo\" and ADD COLUMN \"bar\".\n"
"\n"
"*/\n"
'CREATE TABLE IF NOT EXISTS "Foo" (id TEXT);\n'
)
assert self._run_rules([(self._NEW, sql)]) == []
def test_a_block_comment_does_not_shift_the_reported_line(self):
sql = "/* filler\nfiller */\n" + 'CREATE TABLE "Foo" (id TEXT);\n'
suite = TestMigrationSQLIdempotency()
with pytest.raises(AssertionError) as failure:
suite.test_create_table_uses_if_not_exists([(self._NEW, sql)])
assert f"{self._NEW}:3:" in str(failure.value)
def test_a_new_migration_with_bare_create_table_fails(self):
assert "test_create_table_uses_if_not_exists" in self._run_rules(
[(self._NEW, 'CREATE TABLE "Foo" (id TEXT);\n')]
)
def test_a_new_migration_with_bare_add_column_fails(self):
assert "test_add_column_uses_if_not_exists" in self._run_rules(
[(self._NEW, 'ALTER TABLE "Foo" ADD COLUMN "bar" TEXT;\n')]
)
def test_the_guarded_forms_pass(self):
sql = (
'CREATE TABLE IF NOT EXISTS "Foo" (id TEXT);\n'
'ALTER TABLE "Foo" ADD COLUMN IF NOT EXISTS "bar" TEXT;\n'
'CREATE INDEX IF NOT EXISTS "Foo_bar_idx" ON "Foo"("bar");\n'
)
assert self._run_rules([(self._NEW, sql)]) == []
def test_a_pre_guard_migration_is_exempt_but_a_new_one_is_not(self):
bare = 'CREATE TABLE "Foo" (id TEXT);\n'
exempt = sorted(_PRE_GUARD_MIGRATIONS)[0]
assert self._run_rules([(exempt, bare)]) == []
assert self._run_rules([(self._NEW, bare)]) != []
def test_every_pre_guard_migration_still_exists_on_disk(self):
present = {name for name, _ in _get_all_migrations()}
missing = _PRE_GUARD_MIGRATIONS - present
assert not missing, f"pre-guard entries naming no migration: {sorted(missing)}"
def test_no_pre_guard_entry_is_already_clean(self):
by_name = dict(_get_all_migrations())
redundant = [
name
for name in sorted(_PRE_GUARD_MIGRATIONS)
if not self._run_rules([(TestMigrationGuardScope._NEW, by_name[name])])
]
assert not redundant, f"these no longer violate and should be removed: {redundant}"