mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix: keep schema reconciliation from fighting a partitioned LiteLLM_SpendLogs (#38452)
* fix: keep schema reconciliation from fighting a partitioned LiteLLM_SpendLogs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: scope partitioned SpendLogs detection to Prisma's target schema Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: default partition detection to Prisma's public schema, not current_schema() Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
6b1844442c
commit
fe87b187c6
6 changed files with 411 additions and 6 deletions
|
|
@ -10,6 +10,11 @@
|
|||
-- partitioned, so existing installs are unaffected until you run this.
|
||||
--
|
||||
-- IMPORTANT
|
||||
-- * After partitioning, `prisma db push` (including the proxy's
|
||||
-- --use_prisma_db_push startup mode) is NOT supported: it tries to rewrite
|
||||
-- the primary key back to ("request_id"), which Postgres rejects on a
|
||||
-- partitioned table. The proxy detects this and exits with guidance.
|
||||
-- Use the default startup path (`prisma migrate deploy`) instead.
|
||||
-- * Test on a staging copy first and take a backup.
|
||||
-- * Postgres cannot convert a populated table to partitioned in place, so this
|
||||
-- renames the old table aside and creates a fresh partitioned table.
|
||||
|
|
|
|||
|
|
@ -40,6 +40,65 @@ def _get_prisma_env() -> dict:
|
|||
|
||||
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
|
||||
|
||||
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
|
||||
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
|
||||
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
|
||||
)
|
||||
_SPEND_LOGS_PK_CLAUSE_RE = re.compile(
|
||||
r'^(?:DROP\s+CONSTRAINT\s+"[^"]*_pkey"'
|
||||
r'|ADD\s+(?:CONSTRAINT\s+"[^"]*"\s+)?PRIMARY\s+KEY\s*\([^)]*\))$',
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
|
||||
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
|
||||
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
|
||||
"reconciles the database against schema.prisma, which declares the unpartitioned "
|
||||
"primary key (\"request_id\"), and Postgres rejects that rewrite with: unique "
|
||||
"constraint on partitioned table must include all partitioning columns. Start the "
|
||||
"proxy without --use_prisma_db_push so it uses `prisma migrate deploy`, which only "
|
||||
"applies shipped migrations and leaves the partitioned primary key alone."
|
||||
)
|
||||
|
||||
|
||||
def _without_sql_comments(statement: str) -> str:
|
||||
return "\n".join(
|
||||
line
|
||||
for line in statement.splitlines()
|
||||
if line.strip() and not line.strip().startswith("--")
|
||||
).strip()
|
||||
|
||||
|
||||
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
|
||||
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
|
||||
if not prefix_match:
|
||||
return statement
|
||||
kept = tuple(
|
||||
clause.strip()
|
||||
for clause in statement[prefix_match.end():].split(",\n")
|
||||
if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause.strip())
|
||||
)
|
||||
if not kept:
|
||||
return None
|
||||
return statement[: prefix_match.end()] + ",\n".join(kept)
|
||||
|
||||
|
||||
def filter_partitioned_spend_logs_diff(diff_sql: str) -> str:
|
||||
"""Drop statements from a `prisma migrate diff` script that fight the
|
||||
SpendLogs partitioning runbook (db_scripts/partition_spend_logs.sql): the
|
||||
primary-key rewrite on "LiteLLM_SpendLogs", which Postgres rejects on a
|
||||
partitioned table, and drops of runbook artifacts such as
|
||||
"LiteLLM_SpendLogs_legacy"."""
|
||||
kept = tuple(
|
||||
filtered
|
||||
for statement in diff_sql.split(";")
|
||||
for bare in (_without_sql_comments(statement),)
|
||||
if bare and not _SPEND_LOGS_ARTIFACT_DROP_RE.match(bare)
|
||||
for filtered in (_without_spend_logs_pk_clauses(bare),)
|
||||
if filtered is not None
|
||||
)
|
||||
return "".join(f"{statement};\n\n" for statement in kept)
|
||||
|
||||
|
||||
def _migration_timestamp(name: str) -> int:
|
||||
"""Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name.
|
||||
|
|
@ -355,7 +414,24 @@ class ProxyExtrasDBManager:
|
|||
return
|
||||
logger.info(f"Migration diff created at {diff_sql_path}")
|
||||
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
filtered_sql = filter_partitioned_spend_logs_diff(
|
||||
diff_sql_path.read_text()
|
||||
)
|
||||
diff_sql_path.write_text(filtered_sql)
|
||||
logger.info(
|
||||
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
|
||||
"rewrite and partitioning artifacts from the drift script"
|
||||
)
|
||||
if not filtered_sql.strip():
|
||||
logger.info("Drift script is empty after filtering; nothing to apply")
|
||||
if not mark_all_applied:
|
||||
return
|
||||
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
|
||||
return
|
||||
|
||||
# 2. Run prisma db execute to apply the migration
|
||||
applied_ok = False
|
||||
try:
|
||||
logger.info("Running prisma db execute to apply the migration diff...")
|
||||
result = subprocess.run(
|
||||
|
|
@ -376,6 +452,7 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
logger.info(f"prisma db execute stdout: {result.stdout}")
|
||||
logger.info("✅ Migration diff applied successfully")
|
||||
applied_ok = True
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(f"Failed to apply migration diff: {e.stderr}")
|
||||
except subprocess.TimeoutExpired:
|
||||
|
|
@ -384,6 +461,16 @@ class ProxyExtrasDBManager:
|
|||
# 3. Mark all migrations as applied
|
||||
if not mark_all_applied:
|
||||
return
|
||||
if not applied_ok:
|
||||
logger.warning(
|
||||
"Drift script failed to apply; NOT marking migrations as "
|
||||
"applied so a later migration run can retry them"
|
||||
)
|
||||
return
|
||||
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
|
||||
|
||||
@staticmethod
|
||||
def _mark_migrations_applied(migrations_dir: str):
|
||||
migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
|
||||
logger.info(f"Resolving {len(migration_names)} migrations")
|
||||
for migration_name in migration_names:
|
||||
|
|
@ -410,6 +497,55 @@ class ProxyExtrasDBManager:
|
|||
f"Failed to resolve migration {migration_name}: {e.stderr}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def spend_logs_is_partitioned() -> bool:
|
||||
"""True when the connected database's LiteLLM_SpendLogs is a
|
||||
partitioned table in Prisma's target schema (the `schema` URL param,
|
||||
falling back to Prisma's default target, public), i.e. the operator
|
||||
ran db_scripts/partition_spend_logs.sql. Returns False when psycopg is
|
||||
unavailable or the database cannot be reached, preserving the
|
||||
pre-existing behavior in those cases."""
|
||||
database_url = os.getenv("DATABASE_URL")
|
||||
if not database_url:
|
||||
return False
|
||||
|
||||
try:
|
||||
import psycopg
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
try:
|
||||
with psycopg.connect(
|
||||
cleaned_url, connect_timeout=10, autocommit=True
|
||||
) as conn:
|
||||
row = conn.execute(
|
||||
"SELECT 1 "
|
||||
"FROM pg_partitioned_table pt "
|
||||
"JOIN pg_class c ON c.oid = pt.partrelid "
|
||||
"JOIN pg_namespace n ON n.oid = c.relnamespace "
|
||||
"WHERE c.relname = 'LiteLLM_SpendLogs' "
|
||||
" AND n.nspname = %s",
|
||||
(
|
||||
ProxyExtrasDBManager._prisma_schema_param(database_url)
|
||||
or "public",
|
||||
),
|
||||
).fetchone()
|
||||
except (psycopg.OperationalError, psycopg.DatabaseError):
|
||||
return False
|
||||
return row is not None
|
||||
|
||||
@staticmethod
|
||||
def _prisma_schema_param(url: str) -> Optional[str]:
|
||||
"""The `schema` query param Prisma uses to pick its target schema,
|
||||
or None when the URL does not set one."""
|
||||
from urllib.parse import urlparse, parse_qsl
|
||||
|
||||
return next(
|
||||
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
|
||||
None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _strip_prisma_query_params(url: str) -> str:
|
||||
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
|
||||
|
|
@ -528,7 +664,8 @@ class ProxyExtrasDBManager:
|
|||
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
|
||||
|
||||
if not use_migrate:
|
||||
# Preserve `prisma db push` path unchanged.
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
|
|
@ -972,6 +1109,8 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
raise
|
||||
else:
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
# Use prisma db push with increased timeout
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
|
|
|
|||
|
|
@ -887,6 +887,22 @@ class PrismaManager:
|
|||
return
|
||||
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
|
||||
|
||||
@staticmethod
|
||||
def _raise_if_partitioned_spend_logs() -> None:
|
||||
"""`prisma db push` rewrites a doc-partitioned LiteLLM_SpendLogs
|
||||
primary key back to ("request_id"), which Postgres rejects. Fail fast
|
||||
with guidance instead of retrying into that raw error. No-op when
|
||||
litellm-proxy-extras is absent."""
|
||||
try:
|
||||
from litellm_proxy_extras.utils import (
|
||||
PARTITIONED_SPEND_LOGS_PUSH_ERROR,
|
||||
ProxyExtrasDBManager,
|
||||
)
|
||||
except ImportError:
|
||||
return
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
|
||||
@staticmethod
|
||||
def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool:
|
||||
"""
|
||||
|
|
@ -921,6 +937,7 @@ class PrismaManager:
|
|||
use_v2_resolver=use_v2_resolver,
|
||||
)
|
||||
else:
|
||||
PrismaManager._raise_if_partitioned_spend_logs()
|
||||
# Use prisma db push with increased timeout
|
||||
subprocess.run(
|
||||
[
|
||||
|
|
|
|||
|
|
@ -1321,10 +1321,10 @@ def run_server(
|
|||
use_v2_resolver=use_v2_migration_resolver,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
# v2 resolver raises on unrecoverable migration errors
|
||||
# (e.g. non-idempotent failures, permission issues).
|
||||
# v1 never raises here, so this only fires when the
|
||||
# operator opted into v2.
|
||||
# Raised on unrecoverable migration errors: the v2
|
||||
# resolver's non-idempotent failures and permission
|
||||
# issues, and any `prisma db push` against a
|
||||
# partitioned LiteLLM_SpendLogs.
|
||||
print(
|
||||
f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m",
|
||||
file=sys.stderr,
|
||||
|
|
|
|||
|
|
@ -12,7 +12,11 @@ sys.path.insert(
|
|||
),
|
||||
)
|
||||
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
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(
|
||||
|
|
@ -475,3 +479,205 @@ class TestMigrationGuardScope:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -215,6 +215,44 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
|
|||
assert applied == [True]
|
||||
|
||||
|
||||
def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch):
|
||||
"""A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the
|
||||
primary key back to ("request_id"), which Postgres rejects; the guard must
|
||||
fail fast with guidance instead of running the push."""
|
||||
from litellm.proxy.db.prisma_client import PrismaManager
|
||||
from litellm_proxy_extras.utils import (
|
||||
PARTITIONED_SPEND_LOGS_PUSH_ERROR,
|
||||
ProxyExtrasDBManager,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
|
||||
)
|
||||
with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached
|
||||
"litellm.proxy.db.prisma_client.subprocess.run"
|
||||
) as mock_run:
|
||||
with pytest.raises(RuntimeError) as err:
|
||||
PrismaManager.setup_database(use_migrate=False)
|
||||
|
||||
assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
|
||||
mock_run.assert_not_called()
|
||||
|
||||
|
||||
def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch):
|
||||
from litellm.proxy.db.prisma_client import PrismaManager
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False)
|
||||
)
|
||||
with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic
|
||||
"litellm.proxy.db.prisma_client.subprocess.run"
|
||||
) as mock_run:
|
||||
assert PrismaManager.setup_database(use_migrate=False) is True
|
||||
|
||||
assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
|
||||
|
||||
|
||||
def _entra_jwt(expires_in_seconds: int) -> str:
|
||||
"""A JWT shaped like a real Entra access token, expiring ``expires_in_seconds`` from now."""
|
||||
import base64
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue