feat(db): opt-in REPLICA IDENTITY FULL after prisma migrations (#35267)

Logical replication consumers need FULL replica identity to reconstruct the
old row of an UPDATE or DELETE, and prisma leaves every table it creates at
the postgres default. Operators had to re-apply the setting by hand after
each migration run.

Setting LITELLM_SET_REPLICA_IDENTITY_FULL now re-asserts it on every LiteLLM
table at the end of a successful migration run, through the prisma CLI so the
dependency-free proxy-extras package stays that way. Tables that are already
FULL are skipped, foreign tables in the same schema are left alone, and a
database that refuses the ALTER is reported rather than failing the run.

Resolves LIT-3022

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

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)

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")

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 == []

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]

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