diff --git a/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py new file mode 100644 index 00000000000..dc92e9dca6a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py @@ -0,0 +1,106 @@ +"""Optional post-migration step that raises Postgres REPLICA IDENTITY to FULL. + +Logical-replication consumers (Neon / lakehouse sync and similar) need FULL +replica identity to reconstruct the old row of an UPDATE or DELETE. Prisma +leaves every table it creates at the Postgres default, so the setting has to be +re-applied by hand after each migration run. Setting +``LITELLM_SET_REPLICA_IDENTITY_FULL`` makes every migration run re-assert it. + +The statement goes through the Prisma CLI rather than a Postgres driver because +``litellm-proxy-extras`` has no runtime dependencies, while the CLI is already +required for the migrations themselves. +""" + +import subprocess +import tempfile +from pathlib import Path + +from litellm_proxy_extras._logging import logger + +REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL" + +REPLICA_IDENTITY_FULL_SQL = r""" +DO $$ +DECLARE + target regclass; +BEGIN + SET LOCAL lock_timeout = '5s'; + FOR target IN + SELECT c.oid::regclass + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE c.relkind = 'r' + AND c.relreplident <> 'f' + AND n.nspname = ANY (current_schemas(false)) + AND c.relname LIKE 'LiteLLM\_%' + LOOP + BEGIN + EXECUTE format('ALTER TABLE %s REPLICA IDENTITY FULL', target); + EXCEPTION WHEN lock_not_available THEN + RAISE WARNING 'REPLICA IDENTITY FULL skipped for %: table busy, retrying next run', target; + END; + END LOOP; +END +$$; +""" + + +def apply_replica_identity_full( + schema_path: str, + prisma_command: str, + prisma_env: dict[str, str], +) -> bool: + """Set REPLICA IDENTITY FULL on every LiteLLM table that is not already FULL. + + Never raises. Replication metadata is not needed to serve requests, so + every failure mode is reported and stepped over rather than taking down a + migration run that already succeeded: a database that refuses the ALTER + (most often because the runtime user does not own the tables), a missing + or unrunnable Prisma CLI, a read-only temp directory, or a timeout. + + Returns True when the statement was applied, False when it failed. + """ + logger.info("Applying REPLICA IDENTITY FULL to LiteLLM tables") + try: + with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir: + sql_path = Path(tmp_dir) / "replica_identity_full.sql" + sql_path.write_text(REPLICA_IDENTITY_FULL_SQL) + subprocess.run( + [ + prisma_command, + "db", + "execute", + "--file", + str(sql_path), + "--schema", + schema_path, + ], + timeout=60, + check=True, + capture_output=True, + text=True, + env=prisma_env, + ) + except subprocess.CalledProcessError as e: + logger.error( + "Failed to set REPLICA IDENTITY FULL. Logical replication " + "consumers may reject updates to these tables. Grant table " + "ownership to the migration user, or apply " + "`ALTER TABLE ... REPLICA IDENTITY FULL` by hand. Error: %s", + e.stderr, + ) + return False + except subprocess.TimeoutExpired: + logger.error("Timed out setting REPLICA IDENTITY FULL on LiteLLM tables") + return False + except OSError as e: + logger.error( + "Could not run the REPLICA IDENTITY FULL statement. Logical " + "replication consumers may reject updates to these tables. " + "Error: %s", + e, + ) + return False + + logger.info("REPLICA IDENTITY FULL applied to LiteLLM tables") + return True diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 369b6561931..af822573322 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -10,6 +10,10 @@ from pathlib import Path from typing import Optional from litellm_proxy_extras._logging import logger +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) def str_to_bool(value: Optional[str]) -> bool: @@ -676,6 +680,39 @@ class ProxyExtrasDBManager: finally: os.chdir(original_dir) + @staticmethod + def apply_replica_identity_full_if_requested() -> bool: + """ + Re-assert REPLICA IDENTITY FULL on LiteLLM's tables when the operator + opted in via LITELLM_SET_REPLICA_IDENTITY_FULL. + + Prisma leaves new tables at the Postgres default, which logical + replication consumers reject, so the setting has to be re-applied after + every migration run rather than once by hand. + + Returns: + bool: True if the setting was applied, False if it was not + requested or could not be applied. + """ + if not str_to_bool(os.getenv(REPLICA_IDENTITY_FULL_ENV_VAR)): + return False + try: + schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" + prisma_command = _get_prisma_command() + prisma_env = _get_prisma_env() + except OSError as e: + logger.error( + "Could not resolve the migrations directory for the REPLICA " + "IDENTITY FULL step, skipping it. Error: %s", + e, + ) + return False + return apply_replica_identity_full( + schema_path=schema_path, + prisma_command=prisma_command, + prisma_env=prisma_env, + ) + @staticmethod def setup_database( use_migrate: bool = False, use_v2_resolver: bool = False @@ -694,6 +731,15 @@ class ProxyExtrasDBManager: Returns: bool: True if setup was successful, False otherwise """ + migrated = ProxyExtrasDBManager._run_migrations( + use_migrate=use_migrate, use_v2_resolver=use_v2_resolver + ) + if migrated: + ProxyExtrasDBManager.apply_replica_identity_full_if_requested() + return migrated + + @staticmethod + def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool: if use_v2_resolver: logger.info("Using v2 migration resolver (--use_v2_migration_resolver)") return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index cc608d6e82c..1e0d8f5e010 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -834,6 +834,21 @@ class PrismaManager: dname = os.path.dirname(os.path.dirname(abspath)) return dname + @staticmethod + def _apply_replica_identity_full_if_requested() -> None: + """ + `prisma db push` bypasses litellm-proxy-extras, so the opt-in + REPLICA IDENTITY FULL step has to be driven from here too. + + litellm-proxy-extras is an optional install, so this is a no-op when it + is absent. + """ + try: + from litellm_proxy_extras.utils import ProxyExtrasDBManager + except ImportError: + return + ProxyExtrasDBManager.apply_replica_identity_full_if_requested() + @staticmethod def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool: """ @@ -880,6 +895,7 @@ class PrismaManager: timeout=60, check=True, ) + PrismaManager._apply_replica_identity_full_if_requested() return True except subprocess.TimeoutExpired: verbose_proxy_logger.warning(f"Attempt {attempt + 1} timed out") diff --git a/tests/proxy_migration_tests/test_replica_identity_full.py b/tests/proxy_migration_tests/test_replica_identity_full.py new file mode 100644 index 00000000000..6a88e6994e9 --- /dev/null +++ b/tests/proxy_migration_tests/test_replica_identity_full.py @@ -0,0 +1,159 @@ +"""Coverage for the opt-in REPLICA IDENTITY FULL post-migration step. + +The DB-backed tests run against the same Postgres the migration suite uses, in +a throwaway schema so they cannot disturb the migrated tables. +""" + +import os +import uuid + +import pytest + +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) +from litellm_proxy_extras.utils import ProxyExtrasDBManager + +psycopg = pytest.importorskip("psycopg") + +requires_db = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + + +def _base_url() -> str: + return os.environ["DATABASE_URL"].split("?")[0] + + +def _replica_identities(schema: str) -> dict: + with psycopg.connect(_base_url(), autocommit=True) as conn: + rows = conn.execute( + "SELECT c.relname, c.relreplident FROM pg_class c " + "JOIN pg_namespace n ON n.oid = c.relnamespace " + "WHERE n.nspname = %s AND c.relkind = 'r'", + (schema,), + ).fetchall() + return dict(rows) + + +@pytest.fixture +def scratch_schema(monkeypatch): + """A schema holding two LiteLLM tables and one foreign table, all at the default.""" + schema = f"replica_identity_{uuid.uuid4().hex[:8]}" + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute( + f'CREATE TABLE "{schema}"."LiteLLM_ScratchTable" (id TEXT PRIMARY KEY, note TEXT)' + ) + conn.execute(f'CREATE TABLE "{schema}"."LiteLLM_ScratchSibling" (id TEXT PRIMARY KEY)') + conn.execute(f'CREATE TABLE "{schema}"."ScratchForeignTable" (id TEXT PRIMARY KEY)') + + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") + yield schema + + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + + +@requires_db +def test_applies_full_to_litellm_tables_only(scratch_schema, monkeypatch): + monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true") + + assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True + + identities = _replica_identities(scratch_schema) + assert identities["LiteLLM_ScratchTable"] == "f" + assert identities["LiteLLM_ScratchSibling"] == "f" + assert identities["ScratchForeignTable"] == "d" + + +@requires_db +def test_a_locked_table_does_not_block_the_others(scratch_schema, monkeypatch): + """ALTER TABLE needs an exclusive lock, so a table busy with a long read has + to be skipped for the next run instead of stalling every other table behind it.""" + monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true") + + with psycopg.connect(_base_url()) as holder: + holder.execute(f'SELECT * FROM "{scratch_schema}"."LiteLLM_ScratchTable"') + assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True + + identities = _replica_identities(scratch_schema) + assert identities["LiteLLM_ScratchTable"] == "d" + assert identities["LiteLLM_ScratchSibling"] == "f" + + +@requires_db +def test_leaves_tables_alone_when_not_requested(scratch_schema, monkeypatch): + monkeypatch.delenv(REPLICA_IDENTITY_FULL_ENV_VAR, raising=False) + + assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is False + assert _replica_identities(scratch_schema)["LiteLLM_ScratchTable"] == "d" + + +@requires_db +def test_is_idempotent_across_runs(scratch_schema, monkeypatch): + monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true") + + assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True + assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True + + assert _replica_identities(scratch_schema)["LiteLLM_ScratchTable"] == "f" + + +@requires_db +def test_reports_failure_without_raising(scratch_schema, monkeypatch): + """A run that cannot execute the statement must not take the migration down.""" + monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true") + monkeypatch.setattr( + ProxyExtrasDBManager, + "_get_prisma_dir", + staticmethod(lambda: "/nonexistent/prisma/dir"), + ) + + assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is False + assert _replica_identities(scratch_schema)["LiteLLM_ScratchTable"] == "d" + + +def test_reports_an_unrunnable_prisma_cli_without_raising(tmp_path): + """A deployment without the Prisma CLI on PATH must still finish its + migration run instead of dying on the optional replication step.""" + assert ( + apply_replica_identity_full( + schema_path=str(tmp_path / "schema.prisma"), + prisma_command=str(tmp_path / "no-such-prisma"), + prisma_env={}, + ) + is False + ) + + +def test_setup_database_applies_after_a_successful_migration_run(monkeypatch): + applied = [] + monkeypatch.setattr( + ProxyExtrasDBManager, "_run_migrations", staticmethod(lambda **kwargs: True) + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "apply_replica_identity_full_if_requested", + staticmethod(lambda: applied.append(True)), + ) + + assert ProxyExtrasDBManager.setup_database(use_migrate=True) is True + assert applied == [True] + + +def test_setup_database_skips_replica_identity_when_migrations_fail(monkeypatch): + applied = [] + monkeypatch.setattr( + ProxyExtrasDBManager, "_run_migrations", staticmethod(lambda **kwargs: False) + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "apply_replica_identity_full_if_requested", + staticmethod(lambda: applied.append(True)), + ) + + assert ProxyExtrasDBManager.setup_database(use_migrate=True) is False + assert applied == [] diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index eeaf726941f..08b873dfc44 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -193,3 +193,25 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client( mock_kill.assert_not_called() assert wrapper._original_prisma is mock_new_prisma mock_new_prisma.connect.assert_awaited_once() + + +def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): + """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call + into the opt-in REPLICA IDENTITY FULL step.""" + from litellm.proxy.db.prisma_client import PrismaManager + from litellm_proxy_extras.replica_identity import REPLICA_IDENTITY_FULL_ENV_VAR + from litellm_proxy_extras.utils import ProxyExtrasDBManager + + monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true") + applied = [] + monkeypatch.setattr( + ProxyExtrasDBManager, + "apply_replica_identity_full_if_requested", + staticmethod(lambda: applied.append(True)), + ) + + with patch("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"] + assert applied == [True] diff --git a/tests/test_litellm/proxy/db/test_replica_identity.py b/tests/test_litellm/proxy/db/test_replica_identity.py new file mode 100644 index 00000000000..ecfc6433ab1 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_replica_identity.py @@ -0,0 +1,85 @@ +"""The opt-in REPLICA IDENTITY FULL step, without a database. + +The behavior against real Postgres is covered by +tests/proxy_migration_tests/test_replica_identity_full.py; these pin the two +things that hold with no database at all: the statement handed to the Prisma +CLI, and the promise that no failure of this optional step escapes into a +migration run that already succeeded. +""" + +import subprocess +from pathlib import Path +from unittest.mock import patch + +import pytest + +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) +from litellm_proxy_extras.utils import ProxyExtrasDBManager + + +def test_hands_the_alter_statement_to_the_prisma_cli(): + captured = {} + + def capture(cmd, **kwargs): + captured["cmd"] = cmd + captured["sql"] = Path(cmd[cmd.index("--file") + 1]).read_text() + return subprocess.CompletedProcess(cmd, 0) + + with patch( + "litellm_proxy_extras.replica_identity.subprocess.run", side_effect=capture + ): + applied = apply_replica_identity_full( + schema_path="/somewhere/schema.prisma", + prisma_command="prisma", + prisma_env={"DATABASE_URL": "postgresql://x/y"}, + ) + + assert applied is True + assert captured["cmd"][:3] == ["prisma", "db", "execute"] + assert captured["cmd"][-2:] == ["--schema", "/somewhere/schema.prisma"] + + sql = captured["sql"] + assert "ALTER TABLE %s REPLICA IDENTITY FULL" in sql + assert r"c.relname LIKE 'LiteLLM\_%'" in sql + assert "c.relreplident <> 'f'" in sql + assert "lock_timeout" in sql + + +@pytest.mark.parametrize( + "failure", + [ + subprocess.CalledProcessError(1, "prisma", stderr="must be owner of table"), + subprocess.TimeoutExpired("prisma", 60), + OSError(2, "No such file or directory"), + PermissionError(13, "Read-only file system"), + ], + ids=["rejected", "timed-out", "cli-missing", "read-only-fs"], +) +def test_every_failure_is_reported_instead_of_raised(failure): + with patch( + "litellm_proxy_extras.replica_identity.subprocess.run", side_effect=failure + ): + assert ( + apply_replica_identity_full( + schema_path="/somewhere/schema.prisma", + prisma_command="prisma", + prisma_env={}, + ) + is False + ) + + +def test_an_unusable_migrations_dir_skips_the_step_instead_of_killing_the_run( + tmp_path, monkeypatch +): + """LITELLM_MIGRATION_DIR makes the step copy the migrations tree before it + can run, and that copy is filesystem work that can fail on its own.""" + blocker = tmp_path / "blocker" + blocker.write_text("not a directory") + monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true") + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(blocker / "migrations")) + + assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is False