wip - through models fields -> attached in queries, accesible from instances, creates in add and queryset create

ormar/relations/alias_manager.py

@@ -1,13 +1,14 @@
 import string
 import uuid
 from random import choices
-from typing import Any, Dict, List, TYPE_CHECKING, Type
+from typing import Any, Dict, List, TYPE_CHECKING, Type, Union
 
 import sqlalchemy
 from sqlalchemy import text
 
 if TYPE_CHECKING:  # pragma: no cover
     from ormar import Model
+    from ormar.models import ModelRow
 
 
 def get_table_alias() -> str:
@@ -133,7 +134,7 @@ class AliasManager:
         return alias
 
     def resolve_relation_alias(
-        self, from_model: Type["Model"], relation_name: str
+        self, from_model: Union[Type["Model"], Type["ModelRow"]], relation_name: str
     ) -> str:
         """
         Given model and relation name returns the alias for this relation.

ormar/relations/querysetproxy.py

@@ -44,6 +44,11 @@ class QuerysetProxy(ormar.QuerySetProtocol):
         ].get_related_name()
         self.related_field = self.relation.to.Meta.model_fields[self.related_field_name]
         self.owner_pk_value = self._owner.pk
+        self.through_model_name = (
+            self.related_field.through.get_name()
+            if self.type_ == ormar.RelationType.MULTIPLE
+            else None
+        )
 
     @property
     def queryset(self) -> "QuerySet":
@@ -99,17 +104,20 @@ class QuerysetProxy(ormar.QuerySetProtocol):
             for item in self.relation.related_models[:]:
                 self.relation.remove(item)
 
-    async def create_through_instance(self, child: "T") -> None:
+    async def create_through_instance(self, child: "T", **kwargs: Any) -> None:
         """
         Crete a through model instance in the database for m2m relations.
 
+        :param kwargs: dict of additional keyword arguments for through instance
+        :type kwargs: Any
         :param child: child model instance
         :type child: Model
         """
         model_cls = self.relation.through
         owner_column = self.related_field.default_target_field_name()  # type: ignore
         child_column = self.related_field.default_source_field_name()  # type: ignore
-        kwargs = {owner_column: self._owner.pk, child_column: child.pk}
+        rel_kwargs = {owner_column: self._owner.pk, child_column: child.pk}
+        final_kwargs = {**rel_kwargs, **kwargs}
         if child.pk is None:
             raise ModelPersistenceError(
                 f"You cannot save {child.get_name()} "
@@ -117,7 +125,7 @@ class QuerysetProxy(ormar.QuerySetProtocol):
                 f"Save the child model first."
             )
         expr = model_cls.Meta.table.insert()
-        expr = expr.values(**kwargs)
+        expr = expr.values(**final_kwargs)
         # print("\n", expr.compile(compile_kwargs={"literal_binds": True}))
         await model_cls.Meta.database.execute(expr)
 
@@ -270,12 +278,13 @@ class QuerysetProxy(ormar.QuerySetProtocol):
         :return: created model
         :rtype: Model
         """
+        through_kwargs = kwargs.pop(self.through_model_name, {})
         if self.type_ == ormar.RelationType.REVERSE:
             kwargs[self.related_field.name] = self._owner
         created = await self.queryset.create(**kwargs)
         self._register_related(created)
         if self.type_ == ormar.RelationType.MULTIPLE:
-            await self.create_through_instance(created)
+            await self.create_through_instance(created, **through_kwargs)
         return created
 
     async def get_or_create(self, **kwargs: Any) -> "Model":

ormar/relations/relation.py

@@ -26,6 +26,7 @@ class RelationType(Enum):
     PRIMARY = 1
     REVERSE = 2
     MULTIPLE = 3
+    THROUGH = 4
 
 
 class Relation:
@@ -128,7 +129,7 @@ class Relation:
         :type child: Model
         """
         relation_name = self.field_name
-        if self._type == RelationType.PRIMARY:
+        if self._type in (RelationType.PRIMARY, RelationType.THROUGH):
             self.related_models = child
             self._owner.__dict__[relation_name] = child
         else:

ormar/relations/relation_manager.py

@@ -1,7 +1,7 @@
 from typing import Dict, List, Optional, Sequence, TYPE_CHECKING, Type, TypeVar, Union
 from weakref import proxy
 
-from ormar.fields import BaseField
+from ormar.fields import BaseField, ThroughField
 from ormar.fields.foreign_key import ForeignKeyField
 from ormar.fields.many_to_many import ManyToManyField
 from ormar.relations.relation import Relation, RelationType
@@ -42,6 +42,8 @@ class RelationsManager:
         """
         if issubclass(field, ManyToManyField):
             return RelationType.MULTIPLE
+        if issubclass(field, ThroughField):
+            return RelationType.THROUGH
         return RelationType.PRIMARY if not field.virtual else RelationType.REVERSE
 
     def _add_relation(self, field: Type[BaseField]) -> None:

ormar/relations/relation_proxy.py

@@ -163,19 +163,21 @@ class RelationProxy(list):
             else:
                 await item.delete()
 
-    async def add(self, item: "Model") -> None:
+    async def add(self, item: "Model", **kwargs: Any) -> None:
         """
         Adds child model to relation.
 
         For ManyToMany relations through instance is automatically created.
 
+        :param kwargs: dict of additional keyword arguments for through instance
+        :type kwargs: Any
         :param item: child to add to relation
         :type item: Model
         """
         relation_name = self.related_field_name
         self._check_if_model_saved()
         if self.type_ == ormar.RelationType.MULTIPLE:
-            await self.queryset_proxy.create_through_instance(item)
+            await self.queryset_proxy.create_through_instance(item, **kwargs)
             setattr(item, relation_name, self._owner)
         else:
             setattr(item, relation_name, self._owner)