add association proxy support

Author: 50Bytes-dev

sqlmodel/__init__.py

@@ -117,6 +117,7 @@
 from .main import Field as Field
 from .main import Relationship as Relationship
 from .main import SQLModel as SQLModel
+from .main import SQLModelConfig as SQLModelConfig
 from .orm.session import Session as Session
 from .sql.expression import all_ as all_
 from .sql.expression import and_ as and_

sqlmodel/_compat.py

@@ -243,9 +243,9 @@ def sqlmodel_table_construct(
         # End SQLModel override
 
         fields_values: Dict[str, Any] = {}
-        defaults: Dict[
-            str, Any
-        ] = {}  # keeping this separate from `fields_values` helps us compute `_fields_set`
+        defaults: Dict[str, Any] = (
+            {}
+        )  # keeping this separate from `fields_values` helps us compute `_fields_set`
         for name, field in cls.model_fields.items():
             if field.alias and field.alias in values:
                 fields_values[name] = values.pop(field.alias)
@@ -289,6 +289,10 @@ def sqlmodel_table_construct(
             value = values.get(key, Undefined)
             if value is not Undefined:
                 setattr(self_instance, key, value)
+        for key in self_instance.__sqlalchemy_association_proxies__:
+            value = values.get(key, Undefined)
+            if value is not Undefined:
+                setattr(self_instance, key, value)
         # End SQLModel override
         return self_instance
 

sqlmodel/main.py

@@ -110,7 +110,6 @@
     hybrid_method,
     ColumnProperty,
     declared_attr,
-    AssociationProxy,
 ]
 
 
@@ -498,6 +497,7 @@ def Relationship(
 class SQLModelMetaclass(ModelMetaclass, DeclarativeMeta):
     __sqlmodel_relationships__: Dict[str, RelationshipInfo]
     __sqlalchemy_constructs__: Dict[str, SQLAlchemyConstruct]
+    __sqlalchemy_association_proxies__: Dict[str, AssociationProxy]
     model_config: SQLModelConfig
     model_fields: Dict[str, FieldInfo]
     __config__: Type[SQLModelConfig]
@@ -527,11 +527,14 @@ def __new__(
     ) -> Any:
         relationships: Dict[str, RelationshipInfo] = {}
         sqlalchemy_constructs: Dict[str, SQLAlchemyConstruct] = {}
+        sqlalchemy_association_proxies: Dict[str, AssociationProxy] = {}
         dict_for_pydantic = {}
         original_annotations = get_annotations(class_dict)
         pydantic_annotations = {}
         relationship_annotations = {}
         for k, v in class_dict.items():
+            if isinstance(v, AssociationProxy):
+                sqlalchemy_association_proxies[k] = v
             if isinstance(v, RelationshipInfo):
                 relationships[k] = v
             elif isinstance(
@@ -558,6 +561,7 @@ def __new__(
             "__sqlmodel_relationships__": relationships,
             "__annotations__": pydantic_annotations,
             "__sqlalchemy_constructs__": sqlalchemy_constructs,
+            "__sqlalchemy_association_proxies__": sqlalchemy_association_proxies,
         }
         # Duplicate logic from Pydantic to filter config kwargs because if they are
         # passed directly including the registry Pydantic will pass them over to the
@@ -584,6 +588,9 @@ def __new__(
         for k, v in sqlalchemy_constructs.items():
             setattr(new_cls, k, v)
 
+        for k, v in sqlalchemy_association_proxies.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
@@ -639,10 +646,13 @@ def __init__(
         # 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:
+            for (
+                association_proxy_name,
+                association_proxy,
+            ) in cls.__sqlalchemy_association_proxies__.items():
+                setattr(cls, association_proxy_name, association_proxy)
+
             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
@@ -860,6 +870,7 @@ class SQLModel(BaseModel, metaclass=SQLModelMetaclass, registry=default_registry
     __slots__ = ("__weakref__",)
     __tablename__: ClassVar[Union[str, Callable[..., str]]]
     __sqlmodel_relationships__: ClassVar[Dict[str, RelationshipProperty[Any]]]
+    __sqlalchemy_association_proxies__: ClassVar[Dict[str, AssociationProxy]]
     __name__: ClassVar[str]
     metadata: ClassVar[MetaData]
     __allow_unmapped__ = True  # https://docs.sqlalchemy.org/en/20/changelog/migration_20.html#migration-20-step-six
@@ -909,9 +920,19 @@ def __setattr__(self, name: str, value: Any) -> None:
             # Set in SQLAlchemy, before Pydantic to trigger events and updates
             if is_table_model_class(self.__class__) and is_instrumented(self, name):  # type: ignore[no-untyped-call]
                 set_attribute(self, name, value)
+            # Set in SQLAlchemy association proxies
+            if (
+                is_table_model_class(self.__class__)
+                and name in self.__sqlalchemy_association_proxies__
+            ):
+                association_proxy = self.__sqlalchemy_association_proxies__[name]
+                association_proxy.__set__(self, value)
             # Set in Pydantic model to trigger possible validation changes, only for
             # non relationship values
-            if name not in self.__sqlmodel_relationships__:
+            if (
+                name not in self.__sqlmodel_relationships__
+                and name not in self.__sqlalchemy_association_proxies__
+            ):
                 super().__setattr__(name, value)
 
     def __repr_args__(self) -> Sequence[Tuple[Optional[str], Any]]: