[Feature] Proxy: opt-in v2 migration resolver

[yuneng-berri]

Relevant issues

Summary

Problem

ProxyExtrasDBManager._resolve_all_migrations generates a Prisma schema diff
between the live DB and the shipped schema.prisma and applies it
unconditionally via prisma db execute. This bypasses every migration's
careful SQL (data backfills, multi-step transforms, safety checks). When two
LiteLLM versions contend for the same DB during a rolling deploy — or a
previous-version container restarts unexpectedly — each restart forces the
schema to its own version, producing a thrashing loop.

Fix

Add an opt-in v2 migration resolver. The existing (v1) behavior is
unchanged to avoid breaking users whose DBs have unusual ledger state from
prior thrashing.

Users who have seen thrashing during rolling deploys can opt in with
--use_v2_migration_resolver. The v2 resolver:

• Runs prisma migrate deploy only.
• Recovers from P3005 (baseline) and idempotent P3009/P3018 errors — same
  as v1.
• Never calls _resolve_all_migrations (the diff-and-force recovery
  that caused the thrash).
• Emits a non-blocking warning if the DB has migrations newer than anything
  this build ships (ahead-of-HEAD). Does not refuse to start — users with
  unusual ledger state from past thrashing can still upgrade.

Startup prints a message in default (v1) mode pointing at the opt-in flag
so operators who are affected know the option exists.

Testing

• 10 unit tests in litellm-proxy-extras/tests/test_setup_database_fail_fast.py
  covering the v2 fail-fast paths, URL query-param stripping (for psycopg),
  the timestamp helpers, and pinning the default: v1 still invokes
  _resolve_all_migrations, v2 must not.
• Validated end-to-end with a local harness covering 14 stable versions
  (v1.79.3-stable through v1.83.7-stable) × 3 scenario types (clean
  upgrade, thrashed _prisma_migrations state, ahead-of-HEAD DBs). 840/840
  tests passed across a 10× stability gate against the fix.

Type

🆕 New Feature
✅ Test

[greptile-apps]

Greptile Summary

This PR adds an opt-in --use_v2_migration_resolver flag that replaces the diff-and-force _resolve_all_migrations recovery with a plain prisma migrate deploy loop, fixing schema thrashing during rolling deploys. The v1 default is preserved for backwards compatibility, and the v2 path correctly wraps _resolve_specific_migration failures as RuntimeError — but the P3005 baseline branch does not wrap _create_baseline_migration, which can escape as a raw CalledProcessError and bypass proxy_cli.py's except RuntimeError guard.

Confidence Score: 4/5

Safe to merge after fixing the unguarded _create_baseline_migration call in the P3005 v2 branch.

All prior review-thread issues (psycopg transaction state, db-push error swallowing, _warn_if_db_ahead_of_head DatabaseError, _resolve_specific_migration CalledProcessError escape) have been addressed in this revision. One new P1 remains: in the v2 P3005 recovery branch, _create_baseline_migration can raise subprocess.CalledProcessError, which propagates past the outer try/finally (no outer except in v2, unlike v1's loop-level handler) and crashes proxy startup with an unhandled exception instead of the clean sys.exit(2).

litellm-proxy-extras/litellm_proxy_extras/utils.py — P3005 branch in _setup_database_v2 (around line 573–578)

Important Files Changed

Filename	Overview
litellm-proxy-extras/litellm_proxy_extras/utils.py	Adds v2 migration resolver; _create_baseline_migration in the P3005 branch can escape as a raw CalledProcessError, bypassing proxy_cli's RuntimeError guard.
litellm-proxy-extras/tests/test_setup_database_fail_fast.py	New test file with 10 unit tests covering v2 paths; no test for P3005 baseline-creation failure leaking CalledProcessError.
litellm/proxy/proxy_cli.py	Adds --use_v2_migration_resolver flag and wraps setup_database in try/except RuntimeError + sys.exit(2); correctly backwards-compatible.
litellm/proxy/db/prisma_client.py	Threads use_v2_resolver through to ProxyExtrasDBManager.setup_database; minimal, correct change.
.circleci/config.yml	Adds separate CI job for v2 resolver with Postgres; excludes v2 test from existing Python version jobs to prevent DB collision; checksum-verified uv download.
tests/local_testing/test_basic_python_version.py	Refactors smoke test into a helper, adds v2 variant; clean, no regression risk.
tests/test_litellm/proxy/test_proxy_cli.py	Updates mock assertions to include use_v2_resolver=False; correctly pins default behavior.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[proxy_cli run_server] --> B{use_v2_migration_resolver?}
    B -- No --> C[PrismaManager.setup_database\nuse_v2_resolver=False]
    B -- Yes --> D[PrismaManager.setup_database\nuse_v2_resolver=True]

    C --> E[ProxyExtrasDBManager.setup_database v1]
    E --> F[prisma migrate deploy]
    F -- success --> G[_resolve_all_migrations\npost-migration sanity check]
    F -- P3005 --> H[_create_baseline_migration\nthen _resolve_all_migrations]
    F -- P3009 idempotent --> I[_resolve_specific_migration\nthen _resolve_all_migrations]

    D --> J[_setup_database_v2]
    J --> K[_warn_if_db_ahead_of_head\ninfo only, never blocks]
    K --> L[prisma migrate deploy loop x4]
    L -- success --> M[return True]
    L -- TimeoutExpired --> N[sleep + retry]
    L -- P3005 --> O[_create_baseline_migration\n⚠️ CalledProcessError not wrapped!]
    O --> P[continue loop]
    L -- P3009 idempotent --> Q[_roll_back_migration\n_resolve_specific_migration\nwrapped as RuntimeError]
    Q --> P
    L -- P3018 permission --> R[raise RuntimeError]
    L -- P3018 idempotent --> Q
    L -- other --> R

    D --> S{RuntimeError?}
    S -- Yes --> T[print error + sys.exit 2]
    S -- No --> U{setup_ok?}
    U -- False + enforce check --> V[sys.exit 1]
    U -- True --> W[startup continues]

Loading

_{Reviews (4): Last reviewed commit: "[Fix] v2 resolver: swallow non-connectio..." | Re-trigger Greptile}

[greptile-apps]

return inside with psycopg.connect() leaves transaction in aborted state

In psycopg3 (unlike psycopg2), with conn: calls conn.commit() on a clean exit (__exit__ gets None for exc_type). After UndefinedTable is raised and caught, the server-side transaction is in an aborted state. When return fires, psycopg3's __exit__ sends COMMIT, which PostgreSQL rejects with ERROR: current transaction is aborted, raising psycopg.errors.InFailedSqlTransaction (DatabaseError). That is not caught by the outer except psycopg.OperationalError, so it propagates uncaught through setup_database and crashes fresh-DB startups with a confusing traceback rather than a clean first-run.

The fix is to roll back the aborted transaction before returning, or use autocommit=True:

except psycopg.errors.UndefinedTable:
    conn.rollback()  # clear aborted transaction before the context-manager commits
    return  # fresh DB, table not yet created

Alternatively, open the connection with autocommit=True:

with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:

Suggested change

	except psycopg.errors.UndefinedTable:
	return # fresh DB, table not yet created
	except psycopg.errors.UndefinedTable:
	conn.rollback() # clear aborted transaction before the context-manager commits
	return # fresh DB, table not yet created

[greptile-apps]

Misleading "persistent timeouts" message after non-timeout loop exhaustion

The loop (for attempt in range(4)) is exhausted not only by TimeoutExpired but also by every P3005/P3009/P3018 idempotent recovery (continue). If four consecutive recovery attempts are made (e.g., P3005 baseline, then three successive P3009 idempotent retries), the function falls through to this message even though no timeout occurred, confusing operators trying to diagnose the failure.

Suggested change

	raise RuntimeError(
	"Database migration failed after 4 attempts (persistent timeouts). "
	"Check database connectivity and load."
	)
	raise RuntimeError(
	"Database migration failed after 4 attempts (all attempts used). "
	"Check database connectivity and load."
	)

[yuneng-berri]

@greptile re review please

[yuneng-berri]

@greptile

[greptile-apps]

db push path in v2 swallows CalledProcessError/TimeoutExpired silently

The db push branch in _setup_database_v2 wraps subprocess.run in a bare try/finally (directory restore only). If Prisma returns a non-zero exit code or the call times out, CalledProcessError / TimeoutExpired propagates straight through to proxy_cli.py, which only catches RuntimeError. The result is an unhandled traceback crash rather than a logged warning and a graceful False return.

The v1 path avoids this by catching both exceptions at the outer loop level (except subprocess.TimeoutExpired / except subprocess.CalledProcessError) and returning False after retries. The PR comment "Preserve prisma db push path unchanged" is therefore incorrect — the error-handling contract changed.

The fix is to either catch these exceptions in the db push branch and return False, or raise a RuntimeError (consistent with the rest of v2):

if not use_migrate:
    original_dir = os.getcwd()
    os.chdir(migrations_dir)
    try:
        subprocess.run(
            [_get_prisma_command(), "db", "push", "--accept-data-loss"],
            timeout=60,
            check=True,
            env=_get_prisma_env(),
        )
        return True
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
        raise RuntimeError(
            f"prisma db push failed.\n\nDetail: {e}"
        ) from e
    finally:
        os.chdir(original_dir)

[codecov]

Codecov Report

❌ Patch coverage is 63.63636% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/proxy_cli.py	66.66%	3 Missing ⚠️
litellm/proxy/db/prisma_client.py	50.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

_warn_if_db_ahead_of_head can crash startup on non-connection DB errors

The outer except psycopg.OperationalError only catches connection-level failures. Any psycopg.DatabaseError raised by conn.execute(...) that is not UndefinedTable (e.g., InsufficientPrivilege / error code 42501 if the runtime DB user lacks SELECT on _prisma_migrations) propagates uncaught, crashing proxy startup — directly contradicting the docstring's "Safe no-op if psycopg isn't installed or DB isn't reachable" guarantee.

Fix: widen the outer catch to psycopg.DatabaseError so all DB-layer errors are swallowed consistently with the intent:

        except (psycopg.OperationalError, psycopg.DatabaseError):
            return

[greptile-apps]

_resolve_specific_migration errors bypass proxy_cli.py's RuntimeError guard

In the P3009 recovery path (and the equivalent P3018 path at lines 624–626), _resolve_specific_migration(name) is called inside an active except subprocess.CalledProcessError as e: handler without its own try/except. If this call raises subprocess.CalledProcessError, Python does not re-enter the same except block — the new exception propagates as a raw CalledProcessError. proxy_cli.py only catches RuntimeError, so this crashes startup with an unhandled traceback instead of the intended graceful sys.exit(2).

Wrap the resolve call:

try:
    ProxyExtrasDBManager._resolve_specific_migration(name)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as resolve_err:
    raise RuntimeError(
        f"Failed to mark migration {name} as applied after idempotent recovery. "
        f"Manual intervention may be required.\n\nDetail: {resolve_err}"
    ) from resolve_err
continue

[greptile-apps]

_create_baseline_migration CalledProcessError escapes RuntimeError guard

_create_baseline_migration (line 170 in the existing code) catches CalledProcessError, logs it, then re-raises it as a raw CalledProcessError. When that happens inside the inner except subprocess.CalledProcessError as e: handler here, Python propagates the new exception outward — past the outer try/finally (which only runs os.chdir), past setup_database, and up to proxy_cli.py. Because proxy_cli.py's handler only catches RuntimeError, the process crashes with an unhandled traceback instead of the clean sys.exit(2) the v2 resolver promises.

The v1 path avoids this because its structure has an outer except subprocess.CalledProcessError loop-level handler (line 932 of utils.py) that re-catches the escaping exception. v2 has no such outer handler — only the finally for os.chdir.

Fix: wrap the call (consistent with how _resolve_specific_migration is already wrapped):

if "P3005" in stderr and "database schema is not empty" in stderr:
    logger.info(
        "Schema exists but no migrations ledger — creating baseline"
    )
    try:
        ProxyExtrasDBManager._create_baseline_migration(schema_path)
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as baseline_err:
        raise RuntimeError(
            "Failed to create baseline migration during P3005 recovery. "
            f"Detail: {baseline_err}"
        ) from baseline_err
    continue