Merge pull request #1 from PaleNeutron/sqlalchemy_polymorphic_support

Add support for SQLAlchemy polymorphic models

Author: budroco

sqlmodel/_compat.py

@@ -22,6 +22,8 @@
 from pydantic import VERSION as P_VERSION
 from pydantic import BaseModel
 from pydantic.fields import FieldInfo
+from sqlalchemy import inspect
+from sqlalchemy.orm import Mapper
 from typing_extensions import Annotated, get_args, get_origin
 
 # Reassign variable to make it reexported for mypy
@@ -66,6 +68,35 @@ def _is_union_type(t: Any) -> bool:
 finish_init: ContextVar[bool] = ContextVar("finish_init", default=True)
 
 
+def set_polymorphic_default_value(
+    self_instance: _TSQLModel,
+    values: Dict[str, Any],
+) -> bool:
+    """By default, when init a model, pydantic will set the polymorphic_on
+    value to field default value. But when inherit a model, the polymorphic_on
+    should be set to polymorphic_identity value by default."""
+    cls = type(self_instance)
+    mapper = inspect(cls)
+    ret = False
+    if isinstance(mapper, Mapper):
+        polymorphic_on = mapper.polymorphic_on
+        if polymorphic_on is not None:
+            polymorphic_property = mapper.get_property_by_column(polymorphic_on)
+            field_info = get_model_fields(cls).get(polymorphic_property.key)
+            if field_info:
+                v = values.get(polymorphic_property.key)
+                # if model is inherited or polymorphic_on is not explicitly set
+                # set the polymorphic_on by default
+                if mapper.inherits or v is None:
+                    setattr(
+                        self_instance,
+                        polymorphic_property.key,
+                        mapper.polymorphic_identity,
+                    )
+                    ret = True
+    return ret
+
+
 @contextmanager
 def partial_init() -> Generator[None, None, None]:
     token = finish_init.set(False)
@@ -312,6 +343,8 @@ def sqlmodel_table_construct(
             if value is not Undefined:
                 setattr(self_instance, key, value)
         # End SQLModel override
+        # Override polymorphic_on default value
+        set_polymorphic_default_value(self_instance, values)
         return self_instance
 
     def sqlmodel_validate(

sqlmodel/main.py

@@ -2,6 +2,7 @@
 
 import ipaddress
 import uuid
+import warnings
 import weakref
 from datetime import date, datetime, time, timedelta
 from decimal import Decimal
@@ -43,9 +44,10 @@
 )
 from sqlalchemy import Enum as sa_Enum
 from sqlalchemy.orm import (
+    InstrumentedAttribute,
     Mapped,
+    MappedColumn,
     RelationshipProperty,
-    declared_attr,
     registry,
     relationship,
 )
@@ -539,7 +541,42 @@ def __new__(
         config_kwargs = {
             key: kwargs[key] for key in kwargs.keys() & allowed_config_kwargs
         }
-        new_cls = super().__new__(cls, name, bases, dict_used, **config_kwargs)
+        is_polymorphic = False
+        if IS_PYDANTIC_V2:
+            base_fields = {}
+            base_annotations = {}
+            for base in bases[::-1]:
+                if issubclass(base, BaseModel):
+                    base_fields.update(get_model_fields(base))
+                    base_annotations.update(base.__annotations__)
+                    if hasattr(base, "__sqlmodel_relationships__"):
+                        for k in base.__sqlmodel_relationships__:
+                            # create a dummy attribute to avoid inherit
+                            # pydantic will treat it as class variables, and will not become fields on model instances
+                            anno = base_annotations.get(k, Any)
+                            if get_origin(anno) is not ClassVar:
+                                dummy_anno = ClassVar[anno]
+                                dict_used["__annotations__"][k] = dummy_anno
+
+                    if hasattr(base, "__tablename__"):
+                        is_polymorphic = True
+            # use base_fields overwriting the ones from the class for inherit
+            # if base is a sqlalchemy model, it's attributes will be an InstrumentedAttribute
+            # thus pydantic will use the value of the attribute as the default value
+            base_annotations.update(dict_used["__annotations__"])
+            dict_used["__annotations__"] = base_annotations
+            base_fields.update(dict_used)
+            dict_used = base_fields
+        # if is_polymorphic, disable pydantic `shadows an attribute` warning
+        if is_polymorphic:
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    message="Field name .+ shadows an attribute in parent.+",
+                )
+                new_cls = super().__new__(cls, name, bases, dict_used, **config_kwargs)
+        else:
+            new_cls = super().__new__(cls, name, bases, dict_used, **config_kwargs)
         new_cls.__annotations__ = {
             **relationship_annotations,
             **pydantic_annotations,
@@ -559,9 +596,22 @@ def get_config(name: str) -> Any:
 
         config_table = get_config("table")
         if config_table is True:
+            # sqlalchemy mark a class as table by check if it has __tablename__ attribute
+            # or if __tablename__ is in __annotations__. Only set __tablename__ if it's
+            # a table model
+            if new_cls.__name__ != "SQLModel" and not hasattr(new_cls, "__tablename__"):
+                setattr(new_cls, "__tablename__", new_cls.__name__.lower())  # noqa: B010
             # If it was passed by kwargs, ensure it's also set in config
             set_config_value(model=new_cls, parameter="table", value=config_table)
             for k, v in get_model_fields(new_cls).items():
+                original_v = getattr(new_cls, k, None)
+                if (
+                    isinstance(original_v, InstrumentedAttribute)
+                    and k not in class_dict
+                ):
+                    # The attribute was already set by SQLAlchemy, don't override it
+                    # Needed for polymorphic models, see #36
+                    continue
                 col = get_column_from_field(v)
                 setattr(new_cls, k, col)
             # Set a config flag to tell FastAPI that this should be read with a field
@@ -595,7 +645,15 @@ def __init__(
         # trying to create a new SQLAlchemy, for a new table, with the same name, that
         # triggers an error
         base_is_table = any(is_table_model_class(base) for base in bases)
-        if is_table_model_class(cls) and not base_is_table:
+        _mapper_args = dict_.get("__mapper_args__", {})
+        polymorphic_identity = _mapper_args.get("polymorphic_identity")
+        polymorphic_abstract = _mapper_args.get("polymorphic_abstract")
+        has_polymorphic = (
+            polymorphic_identity is not None or polymorphic_abstract is not None
+        )
+
+        # allow polymorphic models inherit from table models
+        if is_table_model_class(cls) and (not base_is_table or has_polymorphic):
             for rel_name, rel_info in cls.__sqlmodel_relationships__.items():
                 if rel_info.sa_relationship:
                     # There's a SQLAlchemy relationship declared, that takes precedence
@@ -703,13 +761,13 @@ def get_sqlalchemy_type(field: Any) -> Any:
     raise ValueError(f"{type_} has no matching SQLAlchemy type")
 
 
-def get_column_from_field(field: Any) -> Column:  # type: ignore
+def get_column_from_field(field: Any) -> Union[Column, MappedColumn]:  # type: ignore
     if IS_PYDANTIC_V2:
         field_info = field
     else:
         field_info = field.field_info
     sa_column = getattr(field_info, "sa_column", Undefined)
-    if isinstance(sa_column, Column):
+    if isinstance(sa_column, Column) or isinstance(sa_column, MappedColumn):
         return sa_column
     sa_type = get_sqlalchemy_type(field)
     primary_key = getattr(field_info, "primary_key", Undefined)
@@ -773,7 +831,6 @@ def get_column_from_field(field: Any) -> Column:  # type: ignore
 class SQLModel(BaseModel, metaclass=SQLModelMetaclass, registry=default_registry):
     # SQLAlchemy needs to set weakref(s), Pydantic will set the other slots values
     __slots__ = ("__weakref__",)
-    __tablename__: ClassVar[Union[str, Callable[..., str]]]
     __sqlmodel_relationships__: ClassVar[Dict[str, RelationshipProperty[Any]]]
     __name__: ClassVar[str]
     metadata: ClassVar[MetaData]
@@ -837,10 +894,6 @@ def __repr_args__(self) -> Sequence[Tuple[Optional[str], Any]]:
             if not (isinstance(k, str) and k.startswith("_sa_"))
         ]
 
-    @declared_attr  # type: ignore
-    def __tablename__(cls) -> str:
-        return cls.__name__.lower()
-
     @classmethod
     def model_validate(  # type: ignore[override]
         cls: Type[_TSQLModel],