scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

Phase 1: introduce SQLModel models alongside Prisma, with parity guard

Browse source 
Foundation for the multi-phase migration of the proxy persistence layer
from Prisma (prisma-client-py==0.11.0) to SQLModel/SQLAlchemy. This change
is non-functional at runtime: nothing imports the new package from the
existing proxy code paths.

Adds:

* litellm/proxy/db/sqlmodel/schema_parser.py
  Pure-Python parser for the subset of Prisma DSL used by schema.prisma.

* litellm/proxy/db/sqlmodel/_generate.py
  Code generator that emits SQLModel class definitions from a parsed
  schema. Produces Black-formatted output. Run manually after schema
  changes:

      uv run python -m litellm.proxy.db.sqlmodel._generate \
          --schema schema.prisma \
          --out litellm/proxy/db/sqlmodel/models.py

* litellm/proxy/db/sqlmodel/models.py
  SQLModel classes for all 64 Prisma models, hand-editable. Composite
  primary keys, @@unique, @@index (incl. map: '...' renames), @@map,
  String[] arrays, BigInt, Json/JSONB, @updatedAt, @default(uuid()/now()),
  and reserved Python attribute names (metadata -> metadata_) are all
  preserved structurally.

* tests/test_litellm/proxy/db/sqlmodel_orm/{test_schema_parser,test_parity}.py
  20 tests covering the parser unit cases, structural parity between
  every Prisma model and its SQLModel class (columns, nullability, type
  category, ARRAY-ness, primary keys, uniques, indexes), and a strict
  guard that the committed models.py is byte-identical to a fresh
  generator run.

* litellm/proxy/db/sqlmodel/README.md
  Phase 1 plan and the explicit out-of-scope items reserved for later
  phases (session abstraction, raw-SQL hotspot port, per-table call site
  migration, replacement of PrismaWrapper / RoutingPrismaWrapper /
  exception classifier, Alembic migrations, Prisma teardown).

Dependencies:

* Adds 'sqlmodel>=0.0.22,<1.0' to the extra_proxy optional group. This
  pulls in SQLAlchemy as a transitive dep. Existing Prisma deps are
  unchanged.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-05-20 16:27:03 +00:00
parent e59e34bed3
commit b0017fb46d
No known key found for this signature in database
10 changed files with 4221 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

116
litellm/proxy/db/sqlmodel/README.md Normal file
View file

@ -0,0 +1,116 @@
# Prisma -> SQLModel migration -- Phase 1
This package is the foundation for migrating `litellm`'s proxy persistence
layer from Prisma (`prisma-client-py==0.11.0`) to SQLModel/SQLAlchemy. The
migration is multi-phase by necessity -- the proxy has ~1,680 Prisma client
call sites across ~147 production files plus ~177 test files, so any
"big-bang" cutover would be unreviewable and unsafe.
## What this Phase ships
| Artefact | Purpose |
|---|---|
| `schema_parser.py` | Tiny pure-Python parser for the subset of Prisma DSL actually used by `schema.prisma`. Used by the parity test and the generator. |
| `_generate.py` | Code generator that emits SQLModel class definitions from a parsed schema. Run it manually after schema changes. |
| `models.py` | SQLModel classes for **all 64 models** in `schema.prisma`. Hand-editable; the generator only seeds the file. |
| `tests/test_litellm/proxy/db/sqlmodel_orm/test_schema_parser.py` | Unit tests for the parser. |
| `tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py` | Parity test: every Prisma model has a matching SQLModel class with matching columns, nullability, primary keys, uniques, and indexes. Also asserts that the committed `models.py` is byte-identical to a fresh generator run. |
**Nothing in this package is wired into the runtime proxy yet.** Importing
the module has no effect on the existing Prisma-backed code paths -- the
generated classes simply sit alongside the Prisma client and are guarded
against drift by CI.
## Why a generator at all?
The schema is the source of truth and changes frequently. Hand-writing 64
SQLModel classes against a 1,378-line Prisma schema invites typos and
silent drift. The generator gives us one well-tested translation rule per
Prisma construct (`@id`, `@@index`, `String[]`, `@updatedAt`, etc.) and the
parity test catches any regression in either the schema or the generator.
When subsequent phases need to add SQLAlchemy-only behaviour (custom
relationships, hybrid properties, `Mapped[...]` annotations, etc.), edit
`models.py` by hand. The generator's output should still load and the
parity test should still pass; if they don't, the schema and the SQLModel
layer have diverged.
## Re-running the generator
```bash
uv run python -m litellm.proxy.db.sqlmodel._generate \
--schema schema.prisma \
--out litellm/proxy/db/sqlmodel/models.py
```
The parity test fails CI if a schema change isn't accompanied by a
regenerated `models.py`.
## Subsequent phases
The work below is the responsibility of follow-up PRs, in roughly this
order. Each phase is independently testable; do not bundle them.
1. **Session abstraction.** Introduce a thin `DBSession` interface that
wraps the existing `prisma_client` today and a SQLAlchemy
`AsyncSession` tomorrow. Land with zero behaviour change. This is the
prerequisite for incrementally swapping call sites.
2. **CI: keep `schema.prisma` and `models.py` in sync.** Add a workflow
that runs the parity test on every PR (the test already exists -- this
step is just enabling it as a required check).
3. **Port the raw-SQL hotspots.** ~288 `query_raw` / `execute_raw` calls
across ~37 files (concentrated in `spend_management_endpoints.py`,
`db/create_views.py`, focus/cloudzero exporters). These are the
easiest call sites to migrate -- the SQL is already there; we just
swap the executor to a SQLAlchemy `session.execute(text(...))`.
4. **Migrate per-table call sites.** ~55 tables touched across ~1,680
Prisma-client call sites. Parallelise by feature area
(keys/teams/users -> spend/logs -> MCP/managed objects ->
adaptive router/workflows). The session abstraction from phase 1 lets
each call site flip independently.
5. **Replace the custom Prisma reliability layer.** The current
`PrismaWrapper` (RDS IAM token rotation), `RoutingPrismaWrapper`
(read/write split), and `PrismaDBExceptionHandler` (~10 distinct
error type classifications) all need SQLAlchemy-native equivalents.
6. **Swap migrations to Alembic.** The current `litellm-proxy-extras`
package bundles 123 Prisma migration files. Establish an Alembic
baseline matching the live schema, with a documented "first-run
after upgrade" path for existing deployments. The 10 / 123
migrations that contain DML need careful translation; the rest are
pure DDL and can be folded into the baseline for fresh installs.
7. **Tear out Prisma.** Remove `prisma==0.11.0` from `pyproject.toml`,
`prisma generate` from all 7 Dockerfiles, the CI workflows that run
it, the 3 `schema.prisma` copies (with their `check-schema-sync` and
`sync-schema` workflows), and the `litellm-proxy-extras` migration
bundle.
## Risks and gotchas surfaced during Phase 1
* **Reserved attribute names.** SQLModel/SQLAlchemy reserve `metadata`
and `registry` on the mapped class. Several Prisma models have a
`metadata Json` column. The generator emits these as Python attribute
`metadata_` while keeping the on-disk column name `metadata` via
`sa_column_kwargs={'name': 'metadata'}`. Migration of call sites must
use `MyTable.metadata_` in Python.
* **`String[]` (Postgres array columns).** Prisma maps `String[]` to a
Postgres `text[]`. The generator uses
`sqlalchemy.dialects.postgresql.ARRAY(Text())`, which is
Postgres-specific. SQLite-backed test environments will need a
separate fixture path -- this is identical to the current Prisma
situation (`prisma-client-py` on SQLite already requires manual JSON
emulation).
* **`Json` columns default-text quoting.** Prisma's `@default("[]")` and
`@default("{}")` emit `'[]'::jsonb` / `'{}'::jsonb` as the Postgres
`DEFAULT`. The generator preserves both the Python `default_factory`
*and* the `server_default` so migrated rows behave identically when
the column is omitted from an INSERT.
* **`@updatedAt`.** Prisma updates the column from the client. The
generator translates this to a SQLAlchemy `onupdate=lambda: ...
utcnow()` so the behaviour persists when ported off Prisma.
* **`cuid()`.** Only `LiteLLM_CronJob.cronjob_id` uses it. The generator
treats it as opaque-string-equivalent to `uuid()` (which is what every
consumer already assumes).
* **Enums.** The single Prisma enum (`JobStatus`) is emitted as a Python
`str`-Enum and the column is stored as `Text` to match what
`prisma-client-py` already does on Postgres. A real `sa.Enum` can be
introduced later if any call site benefits.

30
litellm/proxy/db/sqlmodel/__init__.py Normal file
View file

@ -0,0 +1,30 @@
"""SQLModel-based ORM definitions for the LiteLLM proxy database.
This package is the foundation for migrating proxy persistence from Prisma to
SQLModel/SQLAlchemy. **Phase 1** (this module's current state) introduces:
* :mod:`schema_parser` -- a small ``schema.prisma`` parser used by the
parity test (and by future code generators).
* :mod:`models` -- hand-maintained, generator-seeded SQLModel classes that
mirror every model in the canonical ``schema.prisma``.
* A parity test (in ``tests/test_litellm/proxy/db/sqlmodel/``) that fails
CI if the SQLModel definitions drift from the Prisma schema.
Nothing in this package is wired into the runtime proxy yet -- importing it
has no effect on existing Prisma-backed code paths. Subsequent phases will:
1. introduce a ``DBSession`` abstraction wrapping Prisma today and SQLAlchemy
tomorrow,
2. port raw-SQL call sites (``query_raw`` / ``execute_raw``) onto SQLAlchemy,
3. migrate per-table call sites (~55 tables) behind the abstraction,
4. rebuild the ``PrismaWrapper`` / ``RoutingPrismaWrapper`` / exception
classifier as SQLAlchemy-native components,
5. swap the migration tool from Prisma to Alembic with a baseline derived
from the current schema state.
See ``litellm/proxy/db/sqlmodel/README.md`` for the full plan.
"""
from litellm.proxy.db.sqlmodel.models import ALL_MODELS
__all__ = ["ALL_MODELS"]

515
litellm/proxy/db/sqlmodel/_generate.py Normal file
View file

@ -0,0 +1,515 @@
"""Generator: emit SQLModel class definitions from ``schema.prisma``.
This is a developer tool, not runtime code. Re-run after schema changes
(or rely on the parity test to flag drift) and copy the output into
:mod:`litellm.proxy.db.sqlmodel.models`. The output is plain Python that
should be reviewed and committed by hand -- this generator is here for
correctness, not for automatic codegen at import time.
Usage::
uv run python -m litellm.proxy.db.sqlmodel._generate \\
--schema schema.prisma \\
--out litellm/proxy/db/sqlmodel/models.py
The generated file is structurally equivalent to ``schema.prisma`` (every
model becomes a SQLModel class with one column per scalar field, plus
table-level constraints and indexes). It does **not** model relations -
those will be added by hand in subsequent migration phases as needed.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Set, Tuple
from litellm.proxy.db.sqlmodel.schema_parser import (
PrismaEnum,
PrismaField,
PrismaModel,
PrismaSchema,
parse_schema_file,
)
# Map Prisma scalar -> (python annotation, SQLAlchemy column type expression).
# We deliberately use SQLAlchemy types (not SQLModel sugar) to match what the
# Prisma migrations have shipped historically: BigInt -> BigInteger,
# String -> Text (Prisma's default for `String` is unbounded text on Postgres),
# Json -> JSONB, etc.
_SCALAR_PY_TYPE = {
"String": "str",
"Int": "int",
"BigInt": "int",
"Float": "float",
"Decimal": "Decimal",
"Boolean": "bool",
"DateTime": "datetime",
"Json": "Any",
"Bytes": "bytes",
}
_SCALAR_SA_TYPE = {
"String": "Text()",
"Int": "Integer()",
"BigInt": "BigInteger()",
"Float": "Double()",
"Decimal": "Numeric()",
"Boolean": "Boolean()",
"DateTime": "DateTime(timezone=True)",
"Json": "JSONB()",
"Bytes": "LargeBinary()",
}
# Default expressions for typical Prisma defaults. Returned strings are
# Python source that produces a SQLAlchemy ``Column(... default=..., server_default=...)``
# argument value. We bias toward server defaults for ``now()``, scalar defaults
# for booleans/numbers, and Python factories for ``uuid()``/``cuid()`` (so the
# value is set at INSERT time, matching prisma-client-py behavior).
def _python_default_for(default_raw: str, base_type: str) -> Optional[str]:
"""Return Python ``default=`` argument source, or ``None``."""
raw = default_raw.strip()
if raw == "uuid()":
return "default_factory=lambda: str(__import__('uuid').uuid4())"
if raw == "cuid()":
# cuid is roughly equivalent to uuid for our purposes; the only
# current user is ``LiteLLM_CronJob.cronjob_id`` and downstream
# consumers treat it as an opaque string.
return "default_factory=lambda: str(__import__('uuid').uuid4())"
if raw == "now()":
return "default_factory=lambda: __import__('datetime').datetime.utcnow()"
if raw == "true":
return "default=True"
if raw == "false":
return "default=False"
if raw == "[]":
return "default_factory=list"
if raw == '"{}"':
return "default_factory=dict"
if raw == '"[]"':
return "default_factory=list"
# Quoted string literal
if raw.startswith('"') and raw.endswith('"'):
return f"default={raw}"
# Numeric literal
try:
float(raw)
return f"default={raw}"
except ValueError:
pass
# Enum reference (e.g. JobStatus value: INACTIVE)
if (
raw.replace("_", "").isalnum()
and raw[:1].isalpha()
and base_type not in _SCALAR_PY_TYPE
):
# we can't reference the Python enum here without an import wrangle,
# so fall back to a string default.
return f'default="{raw}"'
return None
def _server_default_for(default_raw: str, base_type: str) -> Optional[str]:
"""Optional ``server_default=`` to match the existing Postgres DDL."""
raw = default_raw.strip()
if raw == "now()":
return "server_default=text('CURRENT_TIMESTAMP')"
if base_type == "Json":
if raw == '"{}"':
return "server_default=text(\"'{}'\")"
if raw == '"[]"':
return "server_default=text(\"'[]'\")"
return None
def _sa_type_for(field: PrismaField, schema: PrismaSchema) -> str:
"""SQLAlchemy column type expression for a Prisma field."""
if field.is_list:
inner = _SCALAR_SA_TYPE.get(field.base_type, "Text()")
return f"ARRAY({inner})"
if field.base_type in _SCALAR_SA_TYPE:
return _SCALAR_SA_TYPE[field.base_type]
if field.base_type in schema.enums:
# Use a plain Text column; we already index/filter these as strings
# everywhere in production and Prisma's enum type is mostly a
# client-side affair. (Subsequent phases can introduce a real
# ``sa.Enum`` if the call sites benefit from it.)
return "Text()"
return "Text()"
def _py_type_for(field: PrismaField, schema: PrismaSchema) -> str:
if field.base_type in _SCALAR_PY_TYPE:
py = _SCALAR_PY_TYPE[field.base_type]
elif field.base_type in schema.enums:
py = "str"
else:
py = "str"
if field.is_list:
py = f"List[{py}]"
if field.is_optional:
py = f"Optional[{py}]"
return py
# Names SQLAlchemy's Declarative API reserves on a mapped class.
# When a Prisma column collides with one of these we emit the Python attribute
# with a trailing underscore but keep the on-disk column name unchanged via
# ``sa_column_kwargs={'name': '...'}``.
_RESERVED_PY_ATTRS: Set[str] = {"metadata", "registry"}
def _format_field(field: PrismaField, schema: PrismaSchema) -> str:
"""Render one ``Foo: <type> = Field(...)`` line for a SQLModel class."""
py_type = _py_type_for(field, schema)
sa_type = _sa_type_for(field, schema)
field_kwargs: List[str] = [f"sa_type={sa_type}"]
sa_column_kwargs: List[str] = []
py_attr_name = field.name
if field.name in _RESERVED_PY_ATTRS:
py_attr_name = f"{field.name}_"
if field.column_name != py_attr_name:
sa_column_kwargs.append(f"'name': {field.column_name!r}")
if field.is_id:
field_kwargs.append("primary_key=True")
if field.is_unique and not field.is_id:
field_kwargs.append("unique=True")
py_default: Optional[str] = None
srv_default: Optional[str] = None
if field.has_default and field.default_raw is not None:
py_default = _python_default_for(field.default_raw, field.base_type)
srv_default = _server_default_for(field.default_raw, field.base_type)
if py_default is not None:
field_kwargs.append(py_default)
elif field.is_optional:
field_kwargs.append("default=None")
elif field.is_list:
field_kwargs.append("default_factory=list")
if srv_default is not None:
# ``server_default`` lives on the SA column, not on the SQLModel Field.
# _server_default_for returns ``server_default=text('...')``; rip the
# value off and stuff it into sa_column_kwargs so SQLModel forwards it.
value = srv_default.split("=", 1)[1]
sa_column_kwargs.append(f"'server_default': {value}")
if field.has_updated_at:
sa_column_kwargs.append(
"'onupdate': lambda: __import__('datetime').datetime.utcnow()"
)
if sa_column_kwargs:
joined = ", ".join(sa_column_kwargs)
field_kwargs.append(f"sa_column_kwargs={{{joined}}}")
field_args = ", ".join(field_kwargs)
return f" {py_attr_name}: {py_type} = Field({field_args})"
def _format_index_args(model: PrismaModel) -> List[str]:
args: List[str] = []
composite_pk: Tuple[str, ...] = (
model.primary_key if len(model.primary_key) > 1 else ()
)
if composite_pk:
cols = ", ".join(repr(c) for c in composite_pk)
args.append(f"PrimaryKeyConstraint({cols})")
for u in model.uniques:
cols = ", ".join(repr(c) for c in u.fields)
args.append(f"UniqueConstraint({cols})")
for idx in model.indexes:
cols = ", ".join(repr(c) for c in idx.fields)
if idx.map_name:
args.append(f"Index({idx.map_name!r}, {cols})")
else:
# Default index name: <table>_<col>_<col>_idx (matches the
# convention Prisma generates so existing DBs stay happy).
default_name = f"{model.table_name}_" + "_".join(idx.fields) + "_idx"
args.append(f"Index({default_name!r}, {cols})")
return args
def _model_class_name(model: PrismaModel) -> str:
"""Map ``LiteLLM_FooTable`` -> ``LiteLLMFooTable`` (CamelCase, no underscores)."""
parts = model.name.split("_")
return "".join(p[:1].upper() + p[1:] for p in parts if p)
def _render_model_class(model: PrismaModel, schema: PrismaSchema) -> str:
cls_name = _model_class_name(model)
lines: List[str] = []
lines.append(f"class {cls_name}(SQLModel, table=True):")
lines.append(f" __tablename__ = {model.table_name!r}")
index_args = _format_index_args(model)
if index_args:
if len(index_args) == 1:
lines.append(f" __table_args__ = ({index_args[0]},)")
else:
lines.append(" __table_args__ = (")
for arg in index_args:
lines.append(f" {arg},")
lines.append(" )")
lines.append("")
for field in model.fields:
lines.append(_format_field(field, schema))
lines.append("")
return "\n".join(lines)
def _render_enum(enum: PrismaEnum) -> str:
lines = [f"class {enum.name}(str, Enum):"]
for v in enum.values:
lines.append(f" {v} = {v!r}")
lines.append("")
return "\n".join(lines)
_DOCSTRING = '''"""SQLModel ORM definitions mirroring ``schema.prisma``.
THIS FILE IS GENERATED by ``litellm.proxy.db.sqlmodel._generate`` but is
CHECKED IN as ordinary Python source. Hand-edits are allowed -- the parity
test in ``tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py`` will
fail CI if structural drift from ``schema.prisma`` is introduced (in
either direction).
Re-generate with::
uv run python -m litellm.proxy.db.sqlmodel._generate \\
--schema schema.prisma \\
--out litellm/proxy/db/sqlmodel/models.py
Phase 1 of the Prisma -> SQLModel migration only ships these definitions;
nothing in the runtime proxy currently imports them. See the package
README for the multi-phase plan.
"""'''
_FOOTER_TEMPLATE = """
ALL_MODELS: List[Type[SQLModel]] = [
{model_lines}
]
"""
# Map Prisma scalar -> (SA import name, sets has_jsonb, sets has_datetime, sets has_decimal, sets has_any)
_SA_IMPORT_FOR_BASE = {
"String": "Text",
"Int": "Integer",
"BigInt": "BigInteger",
"Float": "Double",
"Decimal": "Numeric",
"Boolean": "Boolean",
"DateTime": "DateTime",
"Json": "JSONB",
"Bytes": "LargeBinary",
}
def _classify_field(
field: PrismaField, schema: PrismaSchema, flags: Dict[str, bool]
) -> Optional[str]:
"""Return the SQLAlchemy import name needed for ``field`` and update ``flags``."""
base = field.base_type
if base in schema.enums:
return "Text"
sa_name = _SA_IMPORT_FOR_BASE.get(base, "Text")
if base == "Json":
flags["jsonb"] = True
flags["any"] = True
elif base == "DateTime":
flags["datetime"] = True
elif base == "Decimal":
flags["decimal"] = True
return sa_name
def _gather_features(schema: PrismaSchema) -> Tuple[Set[str], Dict[str, bool]]:
"""Walk the schema once and return (sqlalchemy import names, feature flags)."""
sa_imports: Set[str] = set()
flags: Dict[str, bool] = {
"optional": False,
"any": False,
"datetime": False,
"decimal": False,
"enum_class": bool(schema.enums),
"indexes": False,
"uniques": False,
"composite_pk": False,
"text_default": False,
"array": False,
"jsonb": False,
}
for model in schema.models.values():
if len(model.primary_key) > 1:
flags["composite_pk"] = True
if model.uniques:
flags["uniques"] = True
if model.indexes:
flags["indexes"] = True
for f in model.fields:
if f.is_optional:
flags["optional"] = True
if f.is_list:
flags["array"] = True
sa = _classify_field(f, schema, flags)
if sa:
sa_imports.add(sa)
if (
f.has_default
and f.default_raw is not None
and _server_default_for(f.default_raw, f.base_type) is not None
):
flags["text_default"] = True
return sa_imports, flags
def _collect_used_symbols(schema: PrismaSchema) -> Set[str]:
"""Return a sentinel-encoded set describing imports needed by the output."""
sa_imports, flags = _gather_features(schema)
if flags["indexes"]:
sa_imports.add("Index")
if flags["composite_pk"]:
sa_imports.add("PrimaryKeyConstraint")
if flags["uniques"]:
sa_imports.add("UniqueConstraint")
if flags["text_default"]:
sa_imports.add("text")
pg_imports: List[str] = []
if flags["array"]:
pg_imports.append("ARRAY")
if flags["jsonb"]:
pg_imports.append("JSONB")
typing_imports: List[str] = ["List", "Type"]
if flags["any"]:
typing_imports.append("Any")
if flags["optional"]:
typing_imports.append("Optional")
stdlib_lines: List[str] = []
if flags["datetime"]:
stdlib_lines.append("from datetime import datetime")
if flags["decimal"]:
stdlib_lines.append("from decimal import Decimal")
if flags["enum_class"]:
stdlib_lines.append("from enum import Enum")
used: Set[str] = set(sa_imports)
used.update(f"_pg::{name}" for name in pg_imports)
used.update(f"_typing::{name}" for name in sorted(set(typing_imports)))
used.update(f"_stdlib::{line}" for line in stdlib_lines)
return used
def _render_imports(schema: PrismaSchema) -> str:
used = _collect_used_symbols(schema)
sa = sorted(
s
for s in used
if not s.startswith("_")
and s
in {
"BigInteger",
"Boolean",
"DateTime",
"Double",
"Index",
"Integer",
"LargeBinary",
"Numeric",
"PrimaryKeyConstraint",
"Text",
"UniqueConstraint",
"text",
}
)
pg = sorted(s.split("::", 1)[1] for s in used if s.startswith("_pg::"))
typing = sorted(s.split("::", 1)[1] for s in used if s.startswith("_typing::"))
stdlib = sorted(s.split("::", 1)[1] for s in used if s.startswith("_stdlib::"))
lines: List[str] = ["from __future__ import annotations", ""]
lines.extend(stdlib)
if stdlib:
lines.append("")
lines.append(f"from typing import {', '.join(typing)}")
lines.append("")
if sa:
if len(sa) == 1:
lines.append(f"from sqlalchemy import {sa[0]}")
else:
lines.append("from sqlalchemy import (")
for s in sa:
lines.append(f" {s},")
lines.append(")")
if pg:
lines.append(f"from sqlalchemy.dialects.postgresql import {', '.join(pg)}")
lines.append("from sqlmodel import Field, SQLModel")
return "\n".join(lines)
def _format_with_black(src: str) -> str:
"""Run Black over ``src`` so generator output matches the committed style.
Black is already a hard CI requirement for this repo (see ``CLAUDE.md``),
so we lean on it as the canonical formatter rather than carrying our own
line-wrapping logic. Falls back to the unformatted source if Black is
unavailable -- the parity test will catch the resulting drift.
"""
try:
import black # type: ignore[import-not-found]
except ImportError:
return src
mode = black.Mode(line_length=88)
try:
return black.format_str(src, mode=mode)
except black.InvalidInput:
return src
def render_module(schema: PrismaSchema) -> str:
"""Render the entire ``models.py`` source for the given schema."""
out: List[str] = [_DOCSTRING, "", _render_imports(schema), ""]
if schema.enums:
for name in sorted(schema.enums):
out.append(_render_enum(schema.enums[name]))
for name in sorted(schema.models):
out.append(_render_model_class(schema.models[name], schema))
model_lines = ",\n".join(
f" {_model_class_name(schema.models[n])}" for n in sorted(schema.models)
)
out.append(_FOOTER_TEMPLATE.format(model_lines=model_lines))
return _format_with_black("\n".join(out))
def main(argv: Optional[Iterable[str]] = None) -> int:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--schema", type=Path, default=Path("schema.prisma"))
parser.add_argument("--out", type=Path, required=True)
args = parser.parse_args(list(argv) if argv is not None else None)
schema = parse_schema_file(args.schema)
src = render_module(schema)
args.out.write_text(src)
sys.stdout.write(
f"wrote {args.out} ({len(schema.models)} models, "
f"{len(schema.enums)} enums)\n"
)
return 0
if __name__ == "__main__": # pragma: no cover
raise SystemExit(main())

2491
litellm/proxy/db/sqlmodel/models.py Normal file
View file

File diff suppressed because it is too large Load diff

473
litellm/proxy/db/sqlmodel/schema_parser.py Normal file
View file

@ -0,0 +1,473 @@
"""Minimal ``schema.prisma`` parser used by the SQLModel parity test.
This is intentionally **not** a full Prisma parser. It targets only the
constructs that actually appear in ``litellm``'s ``schema.prisma`` (as of the
start of the Prisma -> SQLModel migration) and is exercised by the parity
test in ``tests/test_litellm/proxy/db/sqlmodel/``.
The parser produces a structured representation that is easy to compare
against the SQLAlchemy ``MetaData`` of the generated SQLModel classes:
* Top-level ``PrismaSchema`` with ``models`` (dict by model name) and
``enums`` (dict by enum name).
* Each ``PrismaModel`` carries its **scalar** fields, primary key,
uniqueness constraints, and indexes.
* Relation fields (``Foo[]`` / ``Foo? @relation(...)``) are recorded
separately in ``relations`` and are explicitly ignored by the column
parity check -- relations are not columns.
The parser is pure-Python (no third-party deps) so it can run in any test
environment and serve as a building block for future code generators.
"""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
# ---------------------------------------------------------------------------
# Public dataclasses
# ---------------------------------------------------------------------------
@dataclass
class PrismaField:
"""A single scalar (or scalar-array) column on a Prisma model."""
name: str # field name as written in schema.prisma
column_name: str # column name on disk (respects @map(...))
base_type: str # e.g. "String", "Int", "BigInt", "DateTime", "Json", "Bytes", "Float", "Boolean", or an enum name
is_optional: bool # True if `?`
is_list: bool # True if `[]`
is_id: bool # True if marked `@id`
is_unique: bool # True if marked `@unique`
has_default: bool
default_raw: Optional[str] # raw text inside `@default(...)`
has_updated_at: bool # True if marked `@updatedAt`
attributes: List[str] = field(default_factory=list) # raw `@...` attributes
@dataclass
class PrismaRelation:
"""A relation field (``Foo[]`` or ``Foo? @relation(...)``) -- not a column."""
name: str
target_model: str
is_optional: bool
is_list: bool
relation_attributes: List[str] = field(default_factory=list)
@dataclass
class PrismaIndex:
"""A ``@@index([...])`` declaration."""
fields: Tuple[str, ...]
map_name: Optional[str] = None
@dataclass
class PrismaUnique:
"""A ``@@unique([...])`` declaration."""
fields: Tuple[str, ...]
@dataclass
class PrismaModel:
"""A Prisma ``model`` block, scalar columns + constraints only."""
name: str
table_name: str # respects ``@@map("...")``; defaults to model name
fields: List[PrismaField] = field(default_factory=list)
relations: List[PrismaRelation] = field(default_factory=list)
primary_key: Tuple[str, ...] = () # field names (not column names)
uniques: List[PrismaUnique] = field(default_factory=list)
indexes: List[PrismaIndex] = field(default_factory=list)
raw_attributes: List[str] = field(default_factory=list)
def field_by_name(self, name: str) -> Optional[PrismaField]:
for f in self.fields:
if f.name == name:
return f
return None
@dataclass
class PrismaEnum:
name: str
values: Tuple[str, ...]
@dataclass
class PrismaSchema:
models: Dict[str, PrismaModel] = field(default_factory=dict)
enums: Dict[str, PrismaEnum] = field(default_factory=dict)
# ---------------------------------------------------------------------------
# Parser
# ---------------------------------------------------------------------------
# Built-in Prisma scalar types we know how to map.
_SCALAR_TYPES = {
"String",
"Int",
"BigInt",
"Float",
"Decimal",
"Boolean",
"DateTime",
"Json",
"Bytes",
}
_MODEL_RE = re.compile(r"^\s*model\s+(\w+)\s*\{\s*$")
_ENUM_RE = re.compile(r"^\s*enum\s+(\w+)\s*\{\s*$")
_DATASOURCE_RE = re.compile(r"^\s*(datasource|generator)\s+\w+\s*\{\s*$")
_TABLE_ATTR_RE = re.compile(r"^\s*@@(\w+)\s*\((.*)\)\s*$")
_TABLE_MAP_RE = re.compile(r"^\s*@@map\s*\(\s*\"([^\"]+)\"\s*\)\s*$")
def _strip_comment(line: str) -> str:
"""Remove a trailing ``// ...`` comment, ignoring `//` inside quotes."""
out: List[str] = []
in_str = False
i = 0
while i < len(line):
ch = line[i]
if ch == '"' and (i == 0 or line[i - 1] != "\\"):
in_str = not in_str
out.append(ch)
i += 1
continue
if not in_str and ch == "/" and i + 1 < len(line) and line[i + 1] == "/":
break
out.append(ch)
i += 1
return "".join(out).rstrip()
def _split_top_level_commas(s: str) -> List[str]:
"""Split a parenthesized argument list on top-level commas only."""
parts: List[str] = []
depth = 0
in_str = False
buf: List[str] = []
for ch in s:
if ch == '"':
in_str = not in_str
buf.append(ch)
elif in_str:
buf.append(ch)
elif ch in "([{":
depth += 1
buf.append(ch)
elif ch in ")]}":
depth -= 1
buf.append(ch)
elif ch == "," and depth == 0:
parts.append("".join(buf).strip())
buf = []
else:
buf.append(ch)
tail = "".join(buf).strip()
if tail:
parts.append(tail)
return parts
def _extract_attributes(rest: str) -> List[str]:
"""Extract ``@foo(...)`` / ``@foo`` attribute substrings from a field tail."""
attrs: List[str] = []
i = 0
while i < len(rest):
if rest[i] == "@":
j = i + 1
while j < len(rest) and (rest[j].isalnum() or rest[j] in "._"):
j += 1
if j < len(rest) and rest[j] == "(":
depth = 1
k = j + 1
in_str = False
while k < len(rest) and depth > 0:
ch = rest[k]
if ch == '"' and rest[k - 1] != "\\":
in_str = not in_str
elif not in_str:
if ch == "(":
depth += 1
elif ch == ")":
depth -= 1
k += 1
attrs.append(rest[i:k])
i = k
continue
attrs.append(rest[i:j])
i = j
continue
i += 1
return attrs
def _parse_default_value(attr: str) -> Optional[str]:
m = re.match(r"^@default\((.*)\)$", attr)
if not m:
return None
return m.group(1).strip()
def _parse_map_value(attr: str) -> Optional[str]:
m = re.match(r"^@map\(\s*\"([^\"]+)\"\s*\)$", attr)
if not m:
return None
return m.group(1)
def _parse_field_line(line: str) -> Optional[Any]:
"""Parse a single field line inside a model block.
Returns either a ``PrismaField``, a ``PrismaRelation``, or ``None`` if the
line is blank/comment-only.
"""
stripped = _strip_comment(line).strip()
if not stripped:
return None
if stripped.startswith("@@"):
return None # handled separately
parts = stripped.split(None, 2)
if len(parts) < 2:
return None
name = parts[0]
type_token = parts[1]
rest = parts[2] if len(parts) == 3 else ""
is_list = type_token.endswith("[]")
if is_list:
base = type_token[:-2]
is_optional = False
elif type_token.endswith("?"):
base = type_token[:-1]
is_optional = True
else:
base = type_token
is_optional = False
attributes = _extract_attributes(rest)
is_relation = base not in _SCALAR_TYPES and any(
a.startswith("@relation") for a in attributes
)
is_relation = is_relation or (
base not in _SCALAR_TYPES and is_list # `Foo[]` back-reference
)
if is_relation:
return PrismaRelation(
name=name,
target_model=base,
is_optional=is_optional,
is_list=is_list,
relation_attributes=attributes,
)
column_name = name
has_default = False
default_raw: Optional[str] = None
has_updated_at = False
is_id = False
is_unique = False
for attr in attributes:
if attr == "@id":
is_id = True
elif attr == "@unique":
is_unique = True
elif attr == "@updatedAt":
has_updated_at = True
elif attr.startswith("@default("):
has_default = True
default_raw = _parse_default_value(attr)
elif attr.startswith("@map("):
mapped = _parse_map_value(attr)
if mapped is not None:
column_name = mapped
return PrismaField(
name=name,
column_name=column_name,
base_type=base,
is_optional=is_optional,
is_list=is_list,
is_id=is_id,
is_unique=is_unique,
has_default=has_default,
default_raw=default_raw,
has_updated_at=has_updated_at,
attributes=attributes,
)
def _parse_field_list(arg: str) -> Tuple[str, ...]:
"""Parse the field list inside ``@@id([...])`` / ``@@index([...])``.
Field expressions like ``checked_at(sort: Desc)`` are reduced to the bare
field name, which is what we need for parity (SQLAlchemy index objects
don't capture sort direction in the simple comparison we do).
"""
m = re.match(r"^\s*\[(.*)\]\s*(?:,\s*map\s*:\s*\"([^\"]+)\")?\s*$", arg)
if not m:
return ()
inner = m.group(1)
pieces = _split_top_level_commas(inner)
out: List[str] = []
for p in pieces:
# strip ``(sort: Desc)`` etc.
bare = re.sub(r"\(.*\)", "", p).strip()
if bare:
out.append(bare)
return tuple(out)
def _parse_index_attr(arg: str) -> PrismaIndex:
map_name = None
m = re.search(r"map\s*:\s*\"([^\"]+)\"", arg)
if m:
map_name = m.group(1)
fields = _parse_field_list(arg)
return PrismaIndex(fields=fields, map_name=map_name)
def parse_schema(text: str) -> PrismaSchema:
"""Parse a ``schema.prisma`` source string."""
schema = PrismaSchema()
lines = text.splitlines()
i = 0
n = len(lines)
while i < n:
line = _strip_comment(lines[i])
m_model = _MODEL_RE.match(line)
m_enum = _ENUM_RE.match(line)
m_ds = _DATASOURCE_RE.match(line)
if m_ds:
i = _skip_block(lines, i)
continue
if m_enum:
name = m_enum.group(1)
values, i = _consume_enum(lines, i + 1)
schema.enums[name] = PrismaEnum(name=name, values=values)
continue
if m_model:
name = m_model.group(1)
model, i = _consume_model(lines, i + 1, name)
schema.models[name] = model
continue
i += 1
return schema
def parse_schema_file(path: Path) -> PrismaSchema:
return parse_schema(Path(path).read_text())
def _skip_block(lines: List[str], i: int) -> int:
"""Skip a balanced ``{ ... }`` block starting at ``lines[i]``."""
depth = 0
while i < len(lines):
for ch in lines[i]:
if ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return i + 1
i += 1
return i
def _consume_enum(lines: List[str], i: int) -> Tuple[Tuple[str, ...], int]:
values: List[str] = []
while i < len(lines):
stripped = _strip_comment(lines[i]).strip()
if stripped == "}":
return tuple(values), i + 1
if stripped:
# one identifier per line
tok = stripped.split()[0]
values.append(tok)
i += 1
return tuple(values), i
def _consume_model(
lines: List[str], i: int, model_name: str
) -> Tuple[PrismaModel, int]:
model = PrismaModel(name=model_name, table_name=model_name)
while i < len(lines):
raw = lines[i]
stripped_no_comment = _strip_comment(raw).strip()
if stripped_no_comment == "}":
i += 1
break
if not stripped_no_comment:
i += 1
continue
# @@map / @@id / @@unique / @@index / other table-level attrs
m_map = _TABLE_MAP_RE.match(raw)
if m_map:
model.table_name = m_map.group(1)
i += 1
continue
m_attr = _TABLE_ATTR_RE.match(raw)
if m_attr:
kind = m_attr.group(1)
arg = m_attr.group(2).strip()
model.raw_attributes.append(stripped_no_comment)
if kind == "id":
model.primary_key = _parse_field_list(arg)
elif kind == "unique":
model.uniques.append(PrismaUnique(fields=_parse_field_list(arg)))
elif kind == "index":
model.indexes.append(_parse_index_attr(arg))
i += 1
continue
parsed = _parse_field_line(raw)
if parsed is None:
i += 1
continue
if isinstance(parsed, PrismaField):
model.fields.append(parsed)
if parsed.is_id and not model.primary_key:
model.primary_key = (parsed.name,)
elif isinstance(parsed, PrismaRelation):
model.relations.append(parsed)
i += 1
return model, i
# ---------------------------------------------------------------------------
# Convenience helpers used by the parity test
# ---------------------------------------------------------------------------
def column_specs_for(model: PrismaModel) -> Dict[str, Dict[str, Any]]:
"""Return a normalized ``{column_name: spec}`` for parity comparison."""
specs: Dict[str, Dict[str, Any]] = {}
for f in model.fields:
specs[f.column_name] = {
"field_name": f.name,
"base_type": f.base_type,
"is_optional": f.is_optional,
"is_list": f.is_list,
"is_id": f.is_id,
"is_unique": f.is_unique,
"has_default": f.has_default,
"has_updated_at": f.has_updated_at,
}
return specs

1
pyproject.toml
View file

@ -76,6 +76,7 @@ extra_proxy = [
"resend==2.23.0",
"redisvl==0.4.1; python_version < '3.14'",
"a2a-sdk==0.3.24",
"sqlmodel>=0.0.22,<1.0",
]
utils = [
# Not in Docker or PyPI proxy extra.

0
tests/test_litellm/proxy/db/sqlmodel_orm/__init__.py Normal file
View file

350
tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py Normal file
View file

@ -0,0 +1,350 @@
"""Parity test: SQLModel definitions must match ``schema.prisma``.
If this test fails, either:
* ``schema.prisma`` was changed and ``litellm/proxy/db/sqlmodel/models.py``
was not regenerated, OR
* ``models.py`` was hand-edited in a way that no longer reflects the Prisma
schema (which is still the source of truth during the migration).
Re-run the generator and commit the diff::
uv run python -m litellm.proxy.db.sqlmodel._generate \\
--schema schema.prisma \\
--out litellm/proxy/db/sqlmodel/models.py
The test only enforces structural parity that matters for behavioural
equivalence at the database layer:
* every Prisma model has exactly one SQLModel class,
* every scalar Prisma field has a column with the same on-disk name and
nullability,
* primary keys, ``@@unique`` and ``@@index`` clauses match,
* table names (``@@map``) match.
It deliberately does *not* check Python attribute names, type granularity
beyond the broad SQL category, default values, or relation back-refs --
those are implementation details of the SQLModel layer that may diverge
once we hand-tune for SQLAlchemy idioms in later phases.
"""
from __future__ import annotations
from pathlib import Path
from typing import Dict, Set, Tuple
import pytest
from sqlalchemy import Index, PrimaryKeyConstraint, Table, UniqueConstraint
from litellm.proxy.db.sqlmodel.models import ALL_MODELS
from litellm.proxy.db.sqlmodel.schema_parser import (
PrismaModel,
PrismaSchema,
parse_schema_file,
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
def _find_repo_root() -> Path:
p = Path(__file__).resolve()
while not (p / "schema.prisma").exists():
if p.parent == p:
raise RuntimeError("schema.prisma not found in any ancestor directory")
p = p.parent
return p
@pytest.fixture(scope="module")
def prisma_schema() -> PrismaSchema:
return parse_schema_file(_find_repo_root() / "schema.prisma")
@pytest.fixture(scope="module")
def sqlmodel_tables() -> Dict[str, Table]:
return {cls.__tablename__: cls.__table__ for cls in ALL_MODELS}
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _index_signatures(table: Table) -> Set[Tuple[str, ...]]:
"""Set of ``(col1, col2, ...)`` tuples from non-unique SQLAlchemy indexes."""
sigs: Set[Tuple[str, ...]] = set()
for ix in table.indexes:
if ix.unique:
continue
sigs.add(tuple(c.name for c in ix.columns))
return sigs
def _unique_signatures(table: Table) -> Set[Tuple[str, ...]]:
sigs: Set[Tuple[str, ...]] = set()
for cons in table.constraints:
if isinstance(cons, UniqueConstraint):
sigs.add(tuple(c.name for c in cons.columns))
for col in table.columns:
if col.unique and not col.primary_key:
sigs.add((col.name,))
return sigs
def _pk_signature(table: Table) -> Tuple[str, ...]:
return tuple(c.name for c in table.primary_key.columns)
def _prisma_pk_columns(model: PrismaModel) -> Tuple[str, ...]:
"""Map field-name PK to column-name PK (respects ``@map``)."""
cols: list[str] = []
for fname in model.primary_key:
f = model.field_by_name(fname)
cols.append(f.column_name if f is not None else fname)
return tuple(cols)
def _prisma_unique_signatures(model: PrismaModel) -> Set[Tuple[str, ...]]:
sigs: Set[Tuple[str, ...]] = set()
for u in model.uniques:
sigs.add(tuple(_field_to_column(model, fn) for fn in u.fields))
for f in model.fields:
if f.is_unique and not f.is_id:
sigs.add((f.column_name,))
return sigs
def _prisma_index_signatures(model: PrismaModel) -> Set[Tuple[str, ...]]:
sigs: Set[Tuple[str, ...]] = set()
for idx in model.indexes:
sigs.add(tuple(_field_to_column(model, fn) for fn in idx.fields))
return sigs
def _field_to_column(model: PrismaModel, fname: str) -> str:
f = model.field_by_name(fname)
return f.column_name if f is not None else fname
# Prisma scalar -> coarse SQL category we expect on the generated column.
_EXPECTED_TYPE_CATEGORIES = {
"String": {"text", "varchar"},
"Int": {"integer"},
"BigInt": {"biginteger", "bigint"},
"Float": {"double", "double_precision", "float"},
"Decimal": {"numeric", "decimal"},
"Boolean": {"boolean"},
"DateTime": {"datetime", "timestamp"},
"Json": {"json", "jsonb"},
"Bytes": {"largebinary", "bytea"},
}
def _column_type_category(col) -> str:
return type(col.type).__name__.lower()
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
def test_one_sqlmodel_class_per_prisma_model(prisma_schema, sqlmodel_tables):
prisma_table_names = {m.table_name for m in prisma_schema.models.values()}
sqlmodel_table_names = set(sqlmodel_tables)
missing_in_sqlmodel = prisma_table_names - sqlmodel_table_names
extra_in_sqlmodel = sqlmodel_table_names - prisma_table_names
assert not missing_in_sqlmodel, (
f"Prisma tables with no SQLModel class: {sorted(missing_in_sqlmodel)}. "
"Did you forget to regenerate models.py?"
)
assert not extra_in_sqlmodel, (
f"SQLModel classes with no Prisma model: {sorted(extra_in_sqlmodel)}. "
"Did you forget to update schema.prisma?"
)
def test_columns_match_for_every_table(prisma_schema, sqlmodel_tables):
failures: list[str] = []
for prisma_model in prisma_schema.models.values():
table = sqlmodel_tables[prisma_model.table_name]
prisma_cols = {f.column_name: f for f in prisma_model.fields}
sqlmodel_cols = {c.name: c for c in table.columns}
missing = set(prisma_cols) - set(sqlmodel_cols)
extra = set(sqlmodel_cols) - set(prisma_cols)
if missing:
failures.append(
f"{prisma_model.table_name}: missing columns in SQLModel: {sorted(missing)}"
)
if extra:
failures.append(
f"{prisma_model.table_name}: unexpected columns in SQLModel: {sorted(extra)}"
)
assert not failures, "\n".join(failures)
def test_column_nullability_matches(prisma_schema, sqlmodel_tables):
failures: list[str] = []
for prisma_model in prisma_schema.models.values():
table = sqlmodel_tables[prisma_model.table_name]
sqlmodel_cols = {c.name: c for c in table.columns}
for f in prisma_model.fields:
col = sqlmodel_cols.get(f.column_name)
if col is None:
continue
expected_nullable = f.is_optional
if col.nullable != expected_nullable:
failures.append(
f"{prisma_model.table_name}.{f.column_name}: "
f"prisma optional={f.is_optional} but SQLModel nullable={col.nullable}"
)
assert not failures, "\n".join(failures)
def test_column_types_in_expected_category(prisma_schema, sqlmodel_tables):
"""Coarse type check: e.g. ``BigInt`` -> a BigInteger-class type, not Integer.
We deliberately do not enforce exact ``server_default`` or precision -- those
are implementation details that can drift without behavioural impact, and
they are guarded separately by the migration tests.
"""
failures: list[str] = []
for prisma_model in prisma_schema.models.values():
table = sqlmodel_tables[prisma_model.table_name]
sqlmodel_cols = {c.name: c for c in table.columns}
for f in prisma_model.fields:
col = sqlmodel_cols.get(f.column_name)
if col is None:
continue
expected = _EXPECTED_TYPE_CATEGORIES.get(f.base_type)
if expected is None:
# enum reference or unknown scalar -> skip
continue
actual_kind = _column_type_category(col)
ok = any(token in actual_kind for token in expected)
# ARRAY columns wrap an inner type; check the item type instead.
if not ok and "array" in actual_kind and f.is_list:
inner = type(col.type.item_type).__name__.lower()
ok = any(token in inner for token in expected)
if not ok:
failures.append(
f"{prisma_model.table_name}.{f.column_name}: "
f"prisma type={f.base_type}{'[]' if f.is_list else ''} "
f"but SQLModel column type is {actual_kind}"
)
assert not failures, "\n".join(failures)
def test_array_columns_match(prisma_schema, sqlmodel_tables):
failures: list[str] = []
for prisma_model in prisma_schema.models.values():
table = sqlmodel_tables[prisma_model.table_name]
sqlmodel_cols = {c.name: c for c in table.columns}
for f in prisma_model.fields:
col = sqlmodel_cols.get(f.column_name)
if col is None:
continue
actual_is_array = "array" in type(col.type).__name__.lower()
if f.is_list != actual_is_array:
failures.append(
f"{prisma_model.table_name}.{f.column_name}: "
f"prisma is_list={f.is_list} but SQLModel ARRAY={actual_is_array}"
)
assert not failures, "\n".join(failures)
def test_primary_keys_match(prisma_schema, sqlmodel_tables):
failures: list[str] = []
for prisma_model in prisma_schema.models.values():
table = sqlmodel_tables[prisma_model.table_name]
prisma_pk = _prisma_pk_columns(prisma_model)
sqlmodel_pk = _pk_signature(table)
if set(prisma_pk) != set(sqlmodel_pk):
failures.append(
f"{prisma_model.table_name}: prisma PK={prisma_pk} but SQLModel PK={sqlmodel_pk}"
)
assert not failures, "\n".join(failures)
def test_unique_constraints_match(prisma_schema, sqlmodel_tables):
failures: list[str] = []
for prisma_model in prisma_schema.models.values():
table = sqlmodel_tables[prisma_model.table_name]
prisma_uniques = _prisma_unique_signatures(prisma_model)
sqlmodel_uniques = _unique_signatures(table)
# Set comparison ignores ordering of the unique-constraint columns,
# which matches what Postgres treats as logically equivalent.
prisma_norm = {tuple(sorted(s)) for s in prisma_uniques}
sqlmodel_norm = {tuple(sorted(s)) for s in sqlmodel_uniques}
missing = prisma_norm - sqlmodel_norm
extra = sqlmodel_norm - prisma_norm
if missing:
failures.append(
f"{prisma_model.table_name}: missing unique constraints in SQLModel: {sorted(missing)}"
)
if extra:
failures.append(
f"{prisma_model.table_name}: unexpected unique constraints in SQLModel: {sorted(extra)}"
)
assert not failures, "\n".join(failures)
def test_indexes_match(prisma_schema, sqlmodel_tables):
failures: list[str] = []
for prisma_model in prisma_schema.models.values():
table = sqlmodel_tables[prisma_model.table_name]
prisma_idx = _prisma_index_signatures(prisma_model)
sqlmodel_idx = _index_signatures(table)
# We compare ordered tuples here because index column order
# affects which queries the index can serve.
missing = prisma_idx - sqlmodel_idx
extra = sqlmodel_idx - prisma_idx
if missing:
failures.append(
f"{prisma_model.table_name}: missing indexes in SQLModel: {sorted(missing)}"
)
if extra:
failures.append(
f"{prisma_model.table_name}: unexpected indexes in SQLModel: {sorted(extra)}"
)
assert not failures, "\n".join(failures)
def test_generator_output_is_committed(tmp_path):
"""Re-run the generator and assert the result matches the checked-in file.
This is the strongest guard: it catches any drift in either the schema
or the generator (or hand-edits to ``models.py`` that don't roundtrip).
"""
from litellm.proxy.db.sqlmodel import _generate
schema = parse_schema_file(_find_repo_root() / "schema.prisma")
expected = _generate.render_module(schema)
actual = (
_find_repo_root() / "litellm" / "proxy" / "db" / "sqlmodel" / "models.py"
).read_text()
if expected != actual:
# Surface a small diff so the failure message is actionable.
import difflib
diff = "\n".join(
difflib.unified_diff(
actual.splitlines(),
expected.splitlines(),
fromfile="models.py (committed)",
tofile="models.py (regenerated)",
lineterm="",
n=3,
)
)
pytest.fail(
"litellm/proxy/db/sqlmodel/models.py is out of sync with "
"schema.prisma. Run:\n"
" uv run python -m litellm.proxy.db.sqlmodel._generate "
"--schema schema.prisma --out litellm/proxy/db/sqlmodel/models.py\n\n"
f"Diff (truncated to first 60 lines):\n{chr(10).join(diff.splitlines()[:60])}"
)

228
tests/test_litellm/proxy/db/sqlmodel_orm/test_schema_parser.py Normal file
View file

@ -0,0 +1,228 @@
"""Unit tests for the ``schema.prisma`` parser.
Run with::
uv run pytest tests/test_litellm/proxy/db/sqlmodel/test_schema_parser.py -vv
"""
from __future__ import annotations
import textwrap
import pytest
from litellm.proxy.db.sqlmodel.schema_parser import (
PrismaField,
PrismaRelation,
parse_schema,
)
def test_parse_simple_model():
src = textwrap.dedent(
"""
model Foo {
id String @id @default(uuid())
name String @unique
}
"""
)
schema = parse_schema(src)
assert "Foo" in schema.models
foo = schema.models["Foo"]
assert foo.table_name == "Foo"
assert foo.primary_key == ("id",)
assert [f.name for f in foo.fields] == ["id", "name"]
assert foo.fields[0].is_id
assert foo.fields[0].has_default
assert foo.fields[0].default_raw == "uuid()"
assert foo.fields[1].is_unique
def test_optional_and_array_fields():
src = textwrap.dedent(
"""
model Foo {
id String @id
tags String[] @default([])
note String?
}
"""
)
foo = parse_schema(src).models["Foo"]
f_tags = foo.field_by_name("tags")
assert f_tags is not None
assert f_tags.is_list and not f_tags.is_optional
assert f_tags.has_default and f_tags.default_raw == "[]"
f_note = foo.field_by_name("note")
assert f_note is not None
assert f_note.is_optional and not f_note.is_list
def test_composite_primary_key_and_index():
src = textwrap.dedent(
"""
model Foo {
a String
b String
c Int @default(0)
@@id([a, b])
@@index([c])
@@unique([a, c])
}
"""
)
foo = parse_schema(src).models["Foo"]
assert foo.primary_key == ("a", "b")
assert len(foo.indexes) == 1
assert foo.indexes[0].fields == ("c",)
assert len(foo.uniques) == 1
assert foo.uniques[0].fields == ("a", "c")
def test_index_with_map_and_sort():
src = textwrap.dedent(
"""
model Foo {
a String @id
b DateTime
c String
@@index([a, b, c(sort: Desc)], map: "Foo_custom_idx")
}
"""
)
foo = parse_schema(src).models["Foo"]
assert len(foo.indexes) == 1
idx = foo.indexes[0]
assert idx.map_name == "Foo_custom_idx"
assert idx.fields == ("a", "b", "c")
def test_at_map_renames_column():
src = textwrap.dedent(
"""
model Foo {
id String @id
created String @map("created_at")
}
"""
)
foo = parse_schema(src).models["Foo"]
f = foo.field_by_name("created")
assert f is not None
assert f.column_name == "created_at"
def test_at_at_map_renames_table():
src = textwrap.dedent(
"""
model Foo {
id String @id
@@map("foo_table")
}
"""
)
foo = parse_schema(src).models["Foo"]
assert foo.table_name == "foo_table"
def test_relations_are_separated_from_fields():
src = textwrap.dedent(
"""
model Bar {
id String @id
}
model Foo {
id String @id
bar_id String?
bar Bar? @relation(fields: [bar_id], references: [id])
many Bar[]
}
"""
)
foo = parse_schema(src).models["Foo"]
field_names = {f.name for f in foo.fields}
rel_names = {r.name for r in foo.relations}
assert field_names == {"id", "bar_id"}
assert rel_names == {"bar", "many"}
rel_bar = next(r for r in foo.relations if r.name == "bar")
assert rel_bar.target_model == "Bar"
assert rel_bar.is_optional and not rel_bar.is_list
def test_enum_parsed():
src = textwrap.dedent(
"""
enum Status {
ACTIVE
INACTIVE
}
model Foo {
id String @id
status Status @default(INACTIVE)
}
"""
)
schema = parse_schema(src)
assert schema.enums["Status"].values == ("ACTIVE", "INACTIVE")
f = schema.models["Foo"].field_by_name("status")
assert f is not None
assert f.base_type == "Status"
assert f.default_raw == "INACTIVE"
def test_handles_trailing_block_comment_on_model_line():
"""``model Foo { // comment`` should still be recognized as a model."""
src = textwrap.dedent(
"""
model Foo { // a trailing comment after the brace
id String @id
}
"""
)
schema = parse_schema(src)
assert "Foo" in schema.models
def test_strip_comment_handles_quoted_double_slash():
"""A `//` inside a quoted string default must not be treated as a comment."""
src = textwrap.dedent(
"""
model Foo {
id String @id
url String @default("https://example.com")
}
"""
)
foo = parse_schema(src).models["Foo"]
f = foo.field_by_name("url")
assert f is not None
assert f.default_raw == '"https://example.com"'
def test_real_schema_round_trip(tmp_path):
"""Parse the actual repository ``schema.prisma`` and assert basic shape.
This is a smoke test -- the deeper structural parity check lives in
``test_parity.py``.
"""
from pathlib import Path
repo_root = Path(__file__).resolve()
while not (repo_root / "schema.prisma").exists():
if repo_root.parent == repo_root:
pytest.skip("schema.prisma not found in any ancestor directory")
repo_root = repo_root.parent
schema = parse_schema((repo_root / "schema.prisma").read_text())
assert len(schema.models) >= 60
assert "LiteLLM_VerificationToken" in schema.models
assert "LiteLLM_TeamMembership" in schema.models
# composite PK on TeamMembership
assert schema.models["LiteLLM_TeamMembership"].primary_key == (
"user_id",
"team_id",
)

18
uv.lock generated
View file

@ -9,7 +9,7 @@ resolution-markers = [
]
[options]
exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values.
exclude-newer = "2026-05-17T16:09:35.533374581Z"
exclude-newer-span = "P3D"
[manifest]
@ -3219,6 +3219,7 @@ extra-proxy = [
{ name = "prisma" },
{ name = "redisvl" },
{ name = "resend" },
{ name = "sqlmodel" },
]
google = [
{ name = "google-cloud-aiplatform" },
@ -3444,6 +3445,7 @@ requires-dist = [
{ name = "sentry-sdk", marker = "extra == 'proxy-runtime'", specifier = "==2.21.0" },
{ name = "soundfile", marker = "extra == 'proxy'", specifier = "==0.12.1" },
{ name = "soundfile", marker = "extra == 'stt-nvidia-riva'", specifier = ">=0.12.1" },
{ name = "sqlmodel", marker = "extra == 'extra-proxy'", specifier = ">=0.0.22,<1.0" },
{ name = "tiktoken", specifier = ">=0.8.0,<1.0" },
{ name = "tokenizers", specifier = ">=0.21.0,<1.0" },
{ name = "uvicorn", marker = "extra == 'proxy'", specifier = "==0.33.0" },
@ -7183,6 +7185,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e5/30/8519fdde58a7bdf155b714359791ad1dc018b47d60269d5d160d311fdc36/sqlalchemy-2.0.49-py3-none-any.whl", hash = "sha256:ec44cfa7ef1a728e88ad41674de50f6db8cfdb3e2af84af86e0041aaf02d43d0", size = 1942158, upload-time = "2026-04-03T16:53:44.135Z" },
]
[[package]]
name = "sqlmodel"
version = "0.0.38"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pydantic" },
{ name = "sqlalchemy" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/64/0d/26ec1329960ea9430131fe63f63a95ea4cb8971d49c891ff7e1f3255421c/sqlmodel-0.0.38.tar.gz", hash = "sha256:d583ec237b14103809f74e8630032bc40ab68cd6b754a610f0813c56911a547b", size = 86710, upload-time = "2026-04-02T21:03:55.571Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/72/c7/10c60af0607ab6fa136264f7f39d205932218516226d38585324ffda705d/sqlmodel-0.0.38-py3-none-any.whl", hash = "sha256:84e3fa990a77395461ded72a6c73173438ce8449d5c1c4d97fbff1b1df692649", size = 27294, upload-time = "2026-04-02T21:03:56.406Z" },
]
[[package]]
name = "sqlparse"
version = "0.5.5"