mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
e34f43328c
ProxyExtrasDBManager.spend_logs_is_partitioned() (#38452) silently returns False when psycopg can't be imported, and psycopg was never added to the extra_proxy install, so every production image lacks it. Schema reconciliation then generates the unfiltered primary-key rewrite against a genuinely partitioned LiteLLM_SpendLogs and Postgres rejects it, exactly the failure the fix was meant to prevent. Ships psycopg via extra_proxy and logs a warning when it's still missing instead of failing silently.
705 lines
30 KiB
Python
705 lines
30 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 (
|
|
PARTITIONED_SPEND_LOGS_PUSH_ERROR,
|
|
ProxyExtrasDBManager,
|
|
filter_partitioned_spend_logs_diff,
|
|
)
|
|
|
|
# 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}"
|
|
|
|
|
|
_PARTITIONED_DRIFT_SQL = """-- AlterTable
|
|
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;
|
|
|
|
-- AlterTable
|
|
ALTER TABLE "LiteLLM_SpendLogs" DROP CONSTRAINT "LiteLLM_SpendLogs_pkey",
|
|
ADD COLUMN "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
ADD COLUMN "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id");
|
|
|
|
-- DropTable
|
|
DROP TABLE "LiteLLM_SpendLogs_legacy";
|
|
"""
|
|
|
|
|
|
class TestPartitionedSpendLogsDriftFilter:
|
|
"""A doc-partitioned LiteLLM_SpendLogs (db_scripts/partition_spend_logs.sql) has a
|
|
composite primary key that schema.prisma cannot express, so `prisma migrate diff`
|
|
emits a primary-key rewrite that Postgres rejects, aborting the whole drift script
|
|
before its legitimate statements run."""
|
|
|
|
def test_pk_rewrite_and_runbook_artifact_drops_are_removed(self):
|
|
filtered = filter_partitioned_spend_logs_diff(_PARTITIONED_DRIFT_SQL)
|
|
assert 'DROP CONSTRAINT "LiteLLM_SpendLogs_pkey"' not in filtered
|
|
assert 'PRIMARY KEY ("request_id")' not in filtered
|
|
assert "LiteLLM_SpendLogs_legacy" not in filtered
|
|
|
|
def test_legitimate_statements_in_the_same_script_are_kept(self):
|
|
filtered = filter_partitioned_spend_logs_diff(_PARTITIONED_DRIFT_SQL)
|
|
assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;' in filtered
|
|
assert 'ADD COLUMN "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP' in filtered
|
|
assert 'ADD COLUMN "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP' in filtered
|
|
assert filtered.count('ALTER TABLE "LiteLLM_SpendLogs"') == 1
|
|
|
|
def test_an_alter_containing_only_the_pk_rewrite_is_dropped_entirely(self):
|
|
sql = (
|
|
'ALTER TABLE "LiteLLM_SpendLogs" DROP CONSTRAINT "LiteLLM_SpendLogs_pkey",\n'
|
|
'ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id");\n'
|
|
)
|
|
assert filter_partitioned_spend_logs_diff(sql).strip() == ""
|
|
|
|
def test_other_tables_pk_changes_are_untouched(self):
|
|
sql = (
|
|
'ALTER TABLE "LiteLLM_TeamTable" DROP CONSTRAINT "LiteLLM_TeamTable_pkey",\n'
|
|
'ADD CONSTRAINT "LiteLLM_TeamTable_pkey" PRIMARY KEY ("team_id");\n'
|
|
)
|
|
filtered = filter_partitioned_spend_logs_diff(sql)
|
|
assert 'DROP CONSTRAINT "LiteLLM_TeamTable_pkey"' in filtered
|
|
assert 'PRIMARY KEY ("team_id")' in filtered
|
|
|
|
|
|
class _FakeCompleted:
|
|
stdout = ""
|
|
stderr = ""
|
|
|
|
|
|
class TestResolveAllMigrationsLedger:
|
|
def _run(self, monkeypatch, tmp_path, partitioned, execute_fails):
|
|
import subprocess as subprocess_module
|
|
|
|
import litellm_proxy_extras.utils as utils_module
|
|
|
|
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db")
|
|
monkeypatch.delenv("DIRECT_URL", raising=False)
|
|
monkeypatch.setattr(
|
|
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: partitioned)
|
|
)
|
|
monkeypatch.setattr(
|
|
ProxyExtrasDBManager,
|
|
"_get_migration_names",
|
|
staticmethod(lambda migrations_dir: ["20250326162113_baseline"]),
|
|
)
|
|
calls = []
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
calls.append(cmd)
|
|
if "diff" in cmd:
|
|
kwargs["stdout"].write(_PARTITIONED_DRIFT_SQL)
|
|
return _FakeCompleted()
|
|
if "execute" in cmd:
|
|
executed_sql = open(cmd[cmd.index("--file") + 1]).read()
|
|
calls.append(("executed_sql", executed_sql))
|
|
if execute_fails:
|
|
raise subprocess_module.CalledProcessError(1, cmd, stderr="boom")
|
|
return _FakeCompleted()
|
|
return _FakeCompleted()
|
|
|
|
monkeypatch.setattr(utils_module.subprocess, "run", fake_run)
|
|
ProxyExtrasDBManager._resolve_all_migrations(str(tmp_path), "schema.prisma")
|
|
return calls
|
|
|
|
def _resolved(self, calls):
|
|
return [c for c in calls if isinstance(c, list) and "resolve" in c]
|
|
|
|
def _executed_sql(self, calls):
|
|
return next(c[1] for c in calls if isinstance(c, tuple) and c[0] == "executed_sql")
|
|
|
|
def test_failed_drift_apply_does_not_mark_migrations_applied(self, monkeypatch, tmp_path):
|
|
calls = self._run(monkeypatch, tmp_path, partitioned=False, execute_fails=True)
|
|
assert self._resolved(calls) == []
|
|
|
|
def test_successful_drift_apply_still_marks_migrations_applied(self, monkeypatch, tmp_path):
|
|
calls = self._run(monkeypatch, tmp_path, partitioned=False, execute_fails=False)
|
|
assert len(self._resolved(calls)) == 1
|
|
|
|
def test_partitioned_spend_logs_gets_the_filtered_drift_script(self, monkeypatch, tmp_path):
|
|
calls = self._run(monkeypatch, tmp_path, partitioned=True, execute_fails=False)
|
|
executed_sql = self._executed_sql(calls)
|
|
assert 'PRIMARY KEY ("request_id")' not in executed_sql
|
|
assert "LiteLLM_SpendLogs_legacy" not in executed_sql
|
|
assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;' in executed_sql
|
|
assert 'ADD COLUMN "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP' in executed_sql
|
|
assert len(self._resolved(calls)) == 1
|
|
|
|
def test_unpartitioned_spend_logs_drift_script_is_untouched(self, monkeypatch, tmp_path):
|
|
calls = self._run(monkeypatch, tmp_path, partitioned=False, execute_fails=False)
|
|
assert self._executed_sql(calls) == _PARTITIONED_DRIFT_SQL
|
|
|
|
|
|
class TestPartitionedSpendLogsPushGuard:
|
|
def _forbid_subprocess(self, monkeypatch):
|
|
import litellm_proxy_extras.utils as utils_module
|
|
|
|
def fail_run(cmd, **kwargs):
|
|
raise AssertionError(f"subprocess.run should not be called, got: {cmd}")
|
|
|
|
monkeypatch.setattr(utils_module.subprocess, "run", fail_run)
|
|
|
|
def test_v1_db_push_fails_fast_with_guidance(self, monkeypatch):
|
|
monkeypatch.setattr(
|
|
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
|
|
)
|
|
self._forbid_subprocess(monkeypatch)
|
|
with pytest.raises(RuntimeError) as err:
|
|
ProxyExtrasDBManager._run_migrations(use_migrate=False, use_v2_resolver=False)
|
|
assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
|
|
|
|
def test_v2_db_push_fails_fast_with_guidance(self, monkeypatch):
|
|
monkeypatch.setattr(
|
|
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
|
|
)
|
|
self._forbid_subprocess(monkeypatch)
|
|
with pytest.raises(RuntimeError) as err:
|
|
ProxyExtrasDBManager._setup_database_v2(use_migrate=False)
|
|
assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
|
|
|
|
|
|
class _FakeCursor:
|
|
def fetchone(self):
|
|
return (1,)
|
|
|
|
|
|
class _FakePsycopgConn:
|
|
def __init__(self, executed):
|
|
self._executed = executed
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def execute(self, query, params):
|
|
self._executed.append((query, params))
|
|
return _FakeCursor()
|
|
|
|
|
|
class TestSpendLogsPartitionDetectionSchemaScope:
|
|
"""A same-named LiteLLM_SpendLogs in another schema must not trip the
|
|
detector: the catalog lookup has to be scoped to Prisma's target schema."""
|
|
|
|
def _detect(self, monkeypatch, database_url):
|
|
import sys
|
|
import types
|
|
|
|
executed = []
|
|
fake_psycopg = types.ModuleType("psycopg")
|
|
fake_psycopg.connect = lambda url, **kwargs: _FakePsycopgConn(executed)
|
|
fake_psycopg.OperationalError = type("OperationalError", (Exception,), {})
|
|
fake_psycopg.DatabaseError = type("DatabaseError", (Exception,), {})
|
|
monkeypatch.setitem(sys.modules, "psycopg", fake_psycopg)
|
|
monkeypatch.setenv("DATABASE_URL", database_url)
|
|
assert ProxyExtrasDBManager.spend_logs_is_partitioned() is True
|
|
return executed[0]
|
|
|
|
def test_lookup_is_scoped_to_the_schema_url_param(self, monkeypatch):
|
|
query, params = self._detect(
|
|
monkeypatch, "postgresql://u:p@localhost:5432/db?schema=tenant_a"
|
|
)
|
|
assert "pg_namespace" in query
|
|
assert "n.nspname = %s" in query
|
|
assert params == ("tenant_a",)
|
|
|
|
def test_lookup_falls_back_to_public_without_a_schema_param(self, monkeypatch):
|
|
query, params = self._detect(monkeypatch, "postgresql://u:p@localhost:5432/db")
|
|
assert "n.nspname = %s" in query
|
|
assert params == ("public",)
|
|
|
|
def test_only_partitioned_relations_match(self, monkeypatch):
|
|
query, _ = self._detect(monkeypatch, "postgresql://u:p@localhost:5432/db")
|
|
assert "pg_partitioned_table" in query
|
|
|
|
|
|
class TestSpendLogsPartitionDetectionMissingPsycopg:
|
|
"""psycopg ships in the `extra_proxy` install, but a stripped-down image
|
|
can still lack it. When it does, detection must fail closed to False
|
|
(never crash the migration path) and say so loudly, because a silent
|
|
False here is what let a genuinely partitioned LiteLLM_SpendLogs hit the
|
|
unfiltered primary-key rewrite in production."""
|
|
|
|
def test_missing_psycopg_returns_false(self, monkeypatch):
|
|
monkeypatch.setitem(sys.modules, "psycopg", None)
|
|
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db")
|
|
assert ProxyExtrasDBManager.spend_logs_is_partitioned() is False
|
|
|
|
def test_missing_psycopg_logs_a_warning(self, monkeypatch, caplog):
|
|
monkeypatch.setitem(sys.modules, "psycopg", None)
|
|
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db")
|
|
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
|
|
ProxyExtrasDBManager.spend_logs_is_partitioned()
|
|
assert any(
|
|
"psycopg is not installed" in record.message for record in caplog.records
|
|
)
|