Big-bang: replace Prisma with SQLAlchemy + Prisma-compat shim

[mateo-berri]

TL;DR

Rip-and-replace migration of the proxy DB layer from prisma-client-py==0.11.0 to SQLAlchemy + asyncpg in one PR, the way you asked.

This is the counterpart to the phased migration in #28366. Unlike that PR, this one is lossy — it's not finished, and it's not safe to merge as-is. It does, however, prove the rip-and-replace shape is technically achievable in one diff and gives you a concrete breakage inventory to compare against the phased plan.

Test surface	Pre-migration	Post-migration	Pass rate
tests/test_litellm/proxy/db/	263	247	93.9%
tests/test_litellm/proxy/management_endpoints/	1289	1288	99.9%
tests/test_litellm/proxy/db/sqlmodel_orm/ (new, this PR)	38	38	100%
Combined	1590	1573	98.9%

The 17 failures are all in tests that exercise Prisma-internal subprocess / watchdog / RoutingPrismaWrapper mechanics that no longer apply.

---

How the rip-and-replace works

The architectural trick that makes this fit in one PR (instead of touching ~1,680 call sites) is a Prisma-compatible shim behind the existing prisma_client.db.<table>.<method>(...) interface. Every existing call site keeps working unchanged.

New module	Lines	Purpose
litellm/proxy/db/sqlmodel/schema_parser.py	473	Pure-Python parser for the subset of Prisma DSL used by schema.prisma. (Carried forward from #28366.)
litellm/proxy/db/sqlmodel/_generate.py	523	Code generator that emits Black-formatted SQLModel classes from a parsed schema. JSONB columns use with_variant(JSON, 'sqlite') so models load on Postgres and SQLite test envs.
litellm/proxy/db/sqlmodel/models.py	2,546	SQLModel classes for all 64 Prisma models (generator-seeded).
litellm/proxy/db/sqlmodel/engine.py	326	LiteLLMDB async engine + sessionmaker, with pluggable RDS IAM token rotation (engine recreate on each tick) and optional read-replica routing.
litellm/proxy/db/sqlmodel/compat.py	1,004	PrismaCompatClient — the actual shim.
litellm/proxy/db/sqlmodel/errors.py	106	Native error classes (PrismaError, UniqueViolationError, …). SQLAlchemy IntegrityError / NoResultFound / ProgrammingError are translated to these.

Shim surface (what works)

• All ~55 dynamic table accessors, derived from models.ALL_MODELS via name.lower()
• find_unique, find_first, find_many, count
• create, update, upsert, delete, update_many, delete_many, create_many
• where= translation: equals, not, in, notIn, contains (+ mode: 'insensitive'), startswith, endswith, lt, lte, gt, gte, has, AND, OR, NOT
• data= translation: scalar values, {"increment": N}, {"decrement": N}, {"set": V}
• order=, take=, skip=
• query_raw(sql, *params) and query_raw(query=sql, *params) — Postgres $1/$2 placeholders are rewritten to SQLAlchemy :p1/:p2
• execute_raw(sql, *params)
• batch_() returning a batcher with await batcher.commit()
• async with db.tx() as tx: with nested tx.batch_() (matches db_spend_update_writer.py pattern)
• connect(), disconnect(), is_connected(), start_token_refresh_task(), stop_token_refresh_task()

Wiring change (the actual rip-and-replace)

PrismaClient.__init__ in litellm/proxy/utils.py no longer constructs a Prisma() instance / PrismaWrapper / RoutingPrismaWrapper. It builds a LiteLLMDB + PrismaCompatClient and assigns the latter to self.db. _get_engine_pid() returns 0 (no subprocess to track).

The 9 production files that imported from prisma are redirected to the native errors module. The prisma.Json(...) marker in key_management_endpoints.py is replaced with an identity shim (SQLAlchemy JSONB accepts dicts directly).

prisma==0.11.0 is dropped from extra_proxy and proxy-dev in pyproject.toml.

A test-only stand-in (tests/_prisma_compat.py + conftest.py) registers prisma and prisma.errors in sys.modules so legacy from prisma.errors import X in test fixtures still resolves to the native error classes — production code is fully ported.

---

Test verification (per surface)

tests/test_litellm/proxy/db/sqlmodel_orm/ (new, all green)

38 tests, all passing:

• test_schema_parser.py (11) — parser unit tests
• test_parity.py (9) — every Prisma model has a matching SQLModel class with matching columns, nullability, type category, ARRAY-ness, primary keys, uniques, and indexes; committed models.py is byte-identical to a fresh generator run
• test_compat.py (18) — end-to-end shim tests against in-memory SQLite covering CRUD, filters, increment, batch_, tx-rollback, query_raw, execute_raw

tests/test_litellm/proxy/db/ (existing) — 247 / 263 passing

The 16 failures all live in two test modules that exercise Prisma-specific reliability mechanics:

test_prisma_self_heal.py — 15 failures (all subprocess / watchdog tests)
  test_attempt_db_reconnect_should_not_leak_lock_on_timeout_race
  test_attempt_db_reconnect_should_set_cooldown_after_attempt
  test_attempt_db_reconnect_should_skip_when_in_cooldown
  test_attempt_db_reconnect_should_skip_when_lock_timeout_expires
  test_attempt_db_reconnect_should_succeed
  test_db_health_watchdog_should_trigger_reconnect_on_db_error
  test_db_health_watchdog_should_trigger_reconnect_on_probe_timeout
  test_db_health_watchdog_start_stop_lifecycle
  test_engine_confirmed_dead_persists_across_failed_heavy_reconnect
  test_get_generic_data_propagates_when_reconnect_fails
  test_get_generic_data_retries_on_transport_error_for_config_table
  test_recreate_prisma_client_kills_old_engine_without_disconnect
  test_run_reconnect_cycle_timeout_should_use_single_overall_budget
  test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget
  test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client

test_routing_prisma_wrapper.py — 1 failure
  test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails

These tests poke at prisma_client.db._original_prisma._engine.process.pid, send SIGTERM to the engine subprocess, and verify the RoutingPrismaWrapper falls back when the reader Prisma client fails to construct. None of that exists with SQLAlchemy. Replacing them with SQLAlchemy-native equivalents (e.g. engine.dispose() + retry tests, two-engine routing tests) is necessary follow-up before this PR is mergeable.

tests/test_litellm/proxy/management_endpoints/ — 1288 / 1289 passing

Single failure: test_rotate_master_key_model_data_valid_for_prisma — a prisma-Json contract test in the master-key rotation flow. Almost certainly trivial to fix once we either keep the no-op prisma.Json shim available there or update the test's expectations to match plain dicts.

---

⚠️ Known gaps that block merge as-is

I am being explicit about every place this PR is incomplete; it should not be merged without addressing each one.

1. include= (relation eager-loading) is silently ignored

The shim logs a warning and returns rows with the relation attribute unset. ~55 production sites pass include={"litellm_budget_table": True} etc. Adding SQLAlchemy relationship() definitions to every SQLModel class plus selectinload() translation in the shim is the single biggest remaining piece of work. All those call sites currently appear to "work" but produce subtly broken results (the relation attribute is None instead of the expected eager-loaded row).

2. select=, group_by, some/every relation filters

• select= (3 production sites) — currently logs a warning and returns the full row. Probably fine in practice but technically a contract violation.
• group_by (4 production sites — concentrated in spend_management_endpoints.py) — not implemented; calls would fail. Prisma's group_by is non-trivial; raw SQL is the easier path.
• some / every relation filters (2 sites) — NotImplementedError. Both should be ported to explicit JOINs via raw SQL.

3. Reliability layer not yet rebuilt

PrismaWrapper / RoutingPrismaWrapper are dead code (kept in litellm/proxy/db/prisma_client.py only because legacy tests import them). The DB watchdog in PrismaClient (engine PID tracking, zombie reaping, reconnect on subprocess death) is no-op'd via _get_engine_pid() returning 0. SQLAlchemy-native equivalents are needed:

• Engine dispose() + recreate on transport-class errors
• Read/write split via two engines + a per-session router (the LiteLLMDB.reader / writer plumbing exists but isn't wired into the shim's read paths)
• 15 of the 17 failing tests above need to be replaced with SQLAlchemy-native equivalents

4. Migrations: Alembic not added

litellm-proxy-extras still bundles 123 Prisma migration files. Production deployments that re-run prisma migrate deploy on startup will continue to work — but new schema changes can't be made because the SQLModel side isn't versioned by Alembic. Required before any future schema change. The Alembic baseline can be auto-generated from SQLModel.metadata; the 10 / 123 migrations that contain DML need careful translation, the rest fold into the baseline.

5. Build / deploy surface unchanged

• 7 Dockerfiles still run prisma generate at build time
• Makefile install-test-deps still runs prisma generate
• .devcontainer/post-create.sh still runs prisma generate
• 5 GitHub workflows still run prisma generate / prisma db push
• CircleCI config still does
• .github/workflows/check-schema-sync.yml and sync-schema.yml are now operating on a dead artifact
• 3 schema.prisma copies (root + litellm/proxy/ + litellm-proxy-extras/) all still exist
• Helm migrations-job.yaml and Terraform modules still document prisma migrate deploy as the migration command

All of this is rote follow-up.

6. Test-only prisma shim is a smell

tests/_prisma_compat.py and the root conftest.py install stub prisma / prisma.errors modules in sys.modules so test files that still do from prisma import errors as prisma_errors keep working. The right fix is to edit those ~12 test files to import from litellm.proxy.db.sqlmodel.errors directly. I left the shim because doing it right would have ballooned the PR further.

7. Subtle behavioural drift

prisma-client-py returned auto-generated Pydantic-ish classes with specific quirks; the shim returns SQLModel instances. Surface compatibility (attribute access, .dict(), .model_dump()) works because SQLModel inherits from Pydantic — but call sites that compare against prisma.models.LiteLLM_* or rely on Prisma-specific kwargs in row constructors will break. The survey didn't find any such call sites, but a real merge needs an integration test pass against a live Postgres.

8. The single management-endpoint test failure (test_rotate_master_key_model_data_valid_for_prisma) is unaddressed

Triaged but not fixed.

---

Comparison with the phased PR (#28366)

Dimension	Phased (#28366)	Big-bang (this PR)
Scope of one PR	SQLModel definitions + parity test	All of the above + engine + shim + rip-and-replace + 9-file import sweep
Lines added	4,221	6,359
Files touched	9	27
Production code change	None	PrismaClient.__init__ + 9 import sites + prisma.Json shim
Runtime risk	Zero (additive only)	High (production query paths now go through new code)
Test pass rate after	100% (all existing tests untouched)	98.9% (17 known failures, all in Prisma-internal mechanics tests)
Mergeable as-is	Yes	No — see "Known gaps" above
Reviewer cognitive load	One small foundation	Engine + shim + rip-and-replace + import sweep + test stub all at once
Gaps the reviewer can find	0	At least 8, listed above
Gaps the reviewer cannot find	0	Anywhere include= is silently no-op'd in production

Both PRs are draft. #28366 is the recommended path; this PR exists to demonstrate that the rip-and-replace shape is technically achievable but produces a strictly worse review experience and worse safety properties.

Pre-Submission checklist

• Tests added in tests/test_litellm/proxy/db/sqlmodel_orm/ (38 tests covering parser, parity, and end-to-end shim).
• uv run pytest tests/test_litellm/proxy/db/sqlmodel_orm/ — 38/38 passing.
• uv run ruff check litellm/proxy/db/sqlmodel/ — clean.
• uv run black --check litellm/proxy/db/sqlmodel/ tests/test_litellm/proxy/db/sqlmodel_orm/ — clean.
• PR scope is "as isolated as possible" — explicitly no, by request.
• make test-unit passes — no (17 failures listed above).
• Greptile review.

Slack Thread

  

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[devin-ai-integration]

🚅 Hi, thanks for the PR! I'm Agent Shin, the automated triage bot for this repository. What's this and why am I getting it?

I read the description against our contribution rubric. Here's how it lined up:

What you got right:

• ✅ Linked a related GitHub issue
• ✅ Clear problem description
• ✅ Expected vs. actual behavior

What's still missing:

• end-to-end QA proof (video, screenshot, or real commands with output)
• non-mocked proof that the change works against the real system

The PR has strong context and a clear problem/expected-vs-actual description, but it does not include acceptable end-to-end QA proof. The only verification mentioned is test output and test-suite pass rates, which do not satisfy the triage requirement.

If the description isn't updated in the next 24 hours, I'll auto-close this PR. That's not us saying we don't care about the change; we want the open-PR list to mirror what a maintainer can act on right now, so contributors don't get lost in a backlog. A closed PR is a soft "park this for later," not a rejection. Take your time; everything below still works after the close.

During the grace period: just update the PR description with the missing pieces. No need to ping me; I'll re-check on the next sweep and skip the auto-close if it now passes. See what counts as QA proof for the full rubric (a linked issue alone isn't enough; it covers context, not proof).

If the PR does get auto-closed in 24 hours, you still have easy recovery paths:

• Comment @agent-shin reconsider after updating the description. I'll re-evaluate and reopen the PR if it now passes.
• Comment @greptileai to request a fresh Greptile review; that still works even after the PR is closed, and a stronger score is one of the signals that lifts the PR back into the queue. So a low Greptile score isn't a blocker either.

Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.

(I'm an LLM, so I'm not infallible. If you think I got this wrong, ping a maintainer; they'll override me.)