Support hybrid_property, column_property, declared_attr

sqlmodel/main.py

@@ -39,7 +39,9 @@
     inspect,
 )
 from sqlalchemy import Enum as sa_Enum
+from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import (
+    ColumnProperty,
     Mapped,
     RelationshipProperty,
     declared_attr,
@@ -83,6 +85,7 @@
 _T = TypeVar("_T")
 NoArgAnyCallable = Callable[[], Any]
 IncEx = Union[Set[int], Set[str], Dict[int, Any], Dict[str, Any], None]
+SQLAlchemyConstruct = Union[hybrid_property, ColumnProperty, declared_attr]
 
 
 def __dataclass_transform__(
@@ -388,6 +391,7 @@ def Relationship(
 @__dataclass_transform__(kw_only_default=True, field_descriptors=(Field, FieldInfo))
 class SQLModelMetaclass(ModelMetaclass, DeclarativeMeta):
     __sqlmodel_relationships__: Dict[str, RelationshipInfo]
+    __sqlalchemy_constructs__: Dict[str, Any]
     model_config: SQLModelConfig
     model_fields: Dict[str, FieldInfo]
     __config__: Type[SQLModelConfig]
@@ -415,13 +419,20 @@ def __new__(
         **kwargs: Any,
     ) -> Any:
         relationships: Dict[str, RelationshipInfo] = {}
+        sqlalchemy_constructs: Dict[str, SQLAlchemyConstruct] = {}
         dict_for_pydantic = {}
         original_annotations = get_annotations(class_dict)
         pydantic_annotations = {}
         relationship_annotations = {}
         for k, v in class_dict.items():
             if isinstance(v, RelationshipInfo):
                 relationships[k] = v
+            elif isinstance(v, hybrid_property):
+                sqlalchemy_constructs[k] = v
+            elif isinstance(v, ColumnProperty):
+                sqlalchemy_constructs[k] = v
+            elif isinstance(v, declared_attr):
+                sqlalchemy_constructs[k] = v
             else:
                 dict_for_pydantic[k] = v
         for k, v in original_annotations.items():
@@ -434,6 +445,7 @@ def __new__(
             "__weakref__": None,
             "__sqlmodel_relationships__": relationships,
             "__annotations__": pydantic_annotations,
+            "__sqlalchemy_constructs__": sqlalchemy_constructs,
         }
         # Duplicate logic from Pydantic to filter config kwargs because if they are
         # passed directly including the registry Pydantic will pass them over to the
@@ -455,6 +467,11 @@ def __new__(
             **new_cls.__annotations__,
         }
 
+        # We did not provide the sqlalchemy constructs to Pydantic's new function above
+        # so that they wouldn't be modified. Instead we set them directly to the class below:
+        for k, v in sqlalchemy_constructs.items():
+            setattr(new_cls, k, v)
+
         def get_config(name: str) -> Any:
             config_class_value = get_config_value(
                 model=new_cls, parameter=name, default=Undefined
@@ -472,6 +489,8 @@ def get_config(name: str) -> Any:
             set_config_value(model=new_cls, parameter="table", value=config_table)
             for k, v in get_model_fields(new_cls).items():
                 col = get_column_from_field(v)
+                if k in sqlalchemy_constructs:
+                    continue
                 setattr(new_cls, k, col)
             # Set a config flag to tell FastAPI that this should be read with a field
             # in orm_mode instead of preemptively converting it to a dict.
@@ -493,6 +512,9 @@ def get_config(name: str) -> Any:
             setattr(new_cls, "_sa_registry", config_registry)  # noqa: B010
             setattr(new_cls, "metadata", config_registry.metadata)  # noqa: B010
             setattr(new_cls, "__abstract__", True)  # noqa: B010
+            setattr(new_cls, "__pydantic_private__", {})  # noqa: B010
+            setattr(new_cls, "__pydantic_extra__", {})  # noqa: B010
+
         return new_cls
 
     # Override SQLAlchemy, allow both SQLAlchemy and plain Pydantic models
@@ -506,6 +528,9 @@ def __init__(
         base_is_table = any(is_table_model_class(base) for base in bases)
         if is_table_model_class(cls) and not base_is_table:
             for rel_name, rel_info in cls.__sqlmodel_relationships__.items():
+                if rel_name in cls.__sqlalchemy_constructs__:
+                    # Skip hybrid properties
+                    continue
                 if rel_info.sa_relationship:
                     # There's a SQLAlchemy relationship declared, that takes precedence
                     # over anything else, use that and continue with the next attribute

tests/test_column_property.py

@@ -0,0 +1,80 @@
+from typing import List, Optional
+
+from sqlalchemy import case, create_engine, func
+from sqlalchemy.orm import column_property, declared_attr
+from sqlmodel import Field, Relationship, Session, SQLModel, select
+
+
+def test_query(clear_sqlmodel):
+    class Item(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+        value: float
+        hero_id: int = Field(foreign_key="hero.id")
+        hero: "Hero" = Relationship(back_populates="items")
+
+    class Hero(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+        name: str
+        items: List[Item] = Relationship(back_populates="hero")
+
+        @declared_attr
+        def total_items(cls):
+            return column_property(cls._total_items_expression())
+
+        @classmethod
+        def _total_items_expression(cls):
+            return (
+                select(func.coalesce(func.sum(Item.value), 0))
+                .where(Item.hero_id == cls.id)
+                .correlate_except(Item)
+                .label("total_items")
+            )
+
+        @declared_attr
+        def status(cls):
+            return column_property(
+                select(
+                    case(
+                        (cls._total_items_expression() > 0, "active"), else_="inactive"
+                    )
+                ).scalar_subquery()
+            )
+
+    hero_1 = Hero(name="Deadpond")
+    hero_2 = Hero(name="Spiderman")
+
+    engine = create_engine("sqlite://")
+
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        session.add(hero_1)
+        session.add(hero_2)
+        session.commit()
+        session.refresh(hero_1)
+        session.refresh(hero_2)
+
+    item_1 = Item(value=1.0, hero_id=hero_1.id)
+    item_2 = Item(value=2.0, hero_id=hero_1.id)
+
+    with Session(engine) as session:
+        session.add(item_1)
+        session.add(item_2)
+        session.commit()
+        session.refresh(item_1)
+        session.refresh(item_2)
+
+    with Session(engine) as session:
+        hero_statement = select(Hero).where(Hero.total_items > 0.0)
+        hero = session.exec(hero_statement).first()
+        assert hero.name == "Deadpond"
+        assert hero.total_items == 3.0
+        assert hero.status == "active"
+
+    with Session(engine) as session:
+        hero_statement = select(Hero).where(
+            Hero.status == "inactive",
+        )
+        hero = session.exec(hero_statement).first()
+        assert hero.name == "Spiderman"
+        assert hero.total_items == 0.0
+        assert hero.status == "inactive"

tests/test_enums.py

@@ -58,7 +58,19 @@ def sqlite_dump(sql: TypeEngine, *args, **kwargs):
 sqlite_engine = create_mock_engine("sqlite://", sqlite_dump)
 
 
+def _reset_metadata():
+    SQLModel.metadata.clear()
+
+    class FlatModel(SQLModel, table=True):
+        id: uuid.UUID = Field(primary_key=True)
+        enum_field: MyEnum1
+
+    class InheritModel(BaseModel, table=True):
+        pass
+
+
 def test_postgres_ddl_sql(capsys):
+    _reset_metadata()
     SQLModel.metadata.create_all(bind=postgres_engine, checkfirst=False)
 
     captured = capsys.readouterr()
@@ -67,6 +79,7 @@ def test_postgres_ddl_sql(capsys):
 
 
 def test_sqlite_ddl_sql(capsys):
+    _reset_metadata()
     SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False)
 
     captured = capsys.readouterr()

tests/test_hybrid_property.py

@@ -0,0 +1,80 @@
+from typing import List, Optional
+
+from sqlalchemy import case, create_engine, func
+from sqlalchemy.ext.hybrid import hybrid_property
+from sqlmodel import Field, Relationship, Session, SQLModel, select
+
+
+def test_query(clear_sqlmodel):
+    class Item(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+        value: float
+        hero_id: int = Field(foreign_key="hero.id")
+        hero: "Hero" = Relationship(back_populates="items")
+
+    class Hero(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+        name: str
+        items: List[Item] = Relationship(back_populates="hero")
+
+        @hybrid_property
+        def total_items(self):
+            return sum([item.value for item in self.items], 0)
+
+        @total_items.inplace.expression
+        @classmethod
+        def _total_items_expression(cls):
+            return (
+                select(func.coalesce(func.sum(Item.value), 0))
+                .where(Item.hero_id == cls.id)
+                .correlate(cls)
+                .label("total_items")
+            )
+
+        @hybrid_property
+        def status(self):
+            return "active" if self.total_items > 0 else "inactive"
+
+        @status.inplace.expression
+        @classmethod
+        def _status_expression(cls):
+            return select(
+                case((cls.total_items > 0, "active"), else_="inactive")
+            ).label("status")
+
+    hero_1 = Hero(name="Deadpond")
+    hero_2 = Hero(name="Spiderman")
+
+    engine = create_engine("sqlite://")
+
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        session.add(hero_1)
+        session.add(hero_2)
+        session.commit()
+        session.refresh(hero_1)
+        session.refresh(hero_2)
+
+    item_1 = Item(value=1.0, hero_id=hero_1.id)
+    item_2 = Item(value=2.0, hero_id=hero_1.id)
+
+    with Session(engine) as session:
+        session.add(item_1)
+        session.add(item_2)
+        session.commit()
+        session.refresh(item_1)
+        session.refresh(item_2)
+
+    with Session(engine) as session:
+        hero_statement = select(Hero).where(Hero.total_items > 0.0)
+        hero = session.exec(hero_statement).first()
+        assert hero.total_items == 3.0
+        assert hero.status == "active"
+
+    with Session(engine) as session:
+        hero_statement = select(Hero).where(
+            Hero.status == "inactive",
+        )
+        hero = session.exec(hero_statement).first()
+        assert hero.total_items == 0.0
+        assert hero.status == "inactive"