From dd3240c43216803f1ce52c29a9bfffb34ede9cbf Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 27 Feb 2025 10:53:08 -0800 Subject: [PATCH] Support `sqlmodel_rebuild` The `@sqlmodel` decorator introduced by `fquery.sqlmodel` moves schema definitions and types closer to the object model (references) instead of foreign_keys. Instead of writing: ``` class Hero(SQLModel, table=True): ... team_id: int | None = Field(default=None, foreign_key="team.id") ``` you could write: ``` @sqlmodel class Hero: ... team: Team | None = foreign_key("team.id") ``` Everything works the same as before. Only the syntax has changed. However, if `Team` also has a foreign key to `Hero`, then one of them will have to use forward declaration so type checkers work. In those cases we need a two pass solution where in the second pass, with the full type definition available, we can generate the correct code. Refactor some of the existing code to introduce `sqlmodel_rebuild()`. --- sqlmodel/main.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/sqlmodel/main.py b/sqlmodel/main.py index 3532e81a8e..6c9771b0dd 100644 --- a/sqlmodel/main.py +++ b/sqlmodel/main.py @@ -585,9 +585,24 @@ def get_config(name: str) -> Any: setattr(new_cls, "__abstract__", True) # noqa: B010 return new_cls - # Override SQLAlchemy, allow both SQLAlchemy and plain Pydantic models def __init__( cls, classname: str, bases: Tuple[type, ...], dict_: Dict[str, Any], **kw: Any + ) -> None: + return cls.__do_init__(bases, dict_, kw) # type: ignore + + def sqlmodel_rebuild(cls) -> None: + reg = cls._sa_registry + # clear any exisiting mappers for the cls + managers = [m for m in reg._managers if m.class_ == cls] + for m in managers: + reg._dispose_manager_and_mapper(m) + del reg._managers[m] + + return cls.__do_init__(cls.__bases__, cls.__dict__, {}) # type: ignore + + # Override SQLAlchemy, allow both SQLAlchemy and plain Pydantic models + def __do_init__( + cls, classname: str, bases: Tuple[type, ...], dict_: Dict[str, Any], **kw: Any ) -> None: # Only one of the base classes (or the current one) should be a table model # this allows FastAPI cloning a SQLModel for the response_model without