Source code for tortoise.fields.relational

from typing import (
    TYPE_CHECKING,
    Any,
    AsyncGenerator,
    Generator,
    Generic,
    Iterator,
    List,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
    overload,
)

from pypika import Table
from typing_extensions import Literal

from tortoise.exceptions import ConfigurationError, NoValuesFetched, OperationalError
from tortoise.fields.base import CASCADE, SET_NULL, Field, OnDelete

if TYPE_CHECKING:  # pragma: nocoverage
    from tortoise.backends.base.client import BaseDBAsyncClient
    from tortoise.models import Model
    from tortoise.queryset import Q, QuerySet

MODEL = TypeVar("MODEL", bound="Model")


class _NoneAwaitable:
    __slots__ = ()

    def __await__(self) -> Generator[None, None, None]:
        yield None

    def __bool__(self) -> bool:
        return False


NoneAwaitable = _NoneAwaitable()



[docs]
class ReverseRelation(Generic[MODEL]):
    """
    Relation container for :func:`.ForeignKeyField`.
    """

    def __init__(
        self,
        remote_model: Type[MODEL],
        relation_field: str,
        instance: "Model",
        from_field: str,
    ) -> None:
        self.remote_model = remote_model
        self.relation_field = relation_field
        self.instance = instance
        self.from_field = from_field
        self._fetched = False
        self._custom_query = False
        self.related_objects: List[MODEL] = []

    @property
    def _query(self) -> "QuerySet[MODEL]":
        if not self.instance._saved_in_db:
            raise OperationalError(
                "This objects hasn't been instanced, call .save() before calling related queries"
            )
        return self.remote_model.filter(
            **{self.relation_field: getattr(self.instance, self.from_field)}
        )

    def __contains__(self, item: Any) -> bool:
        self._raise_if_not_fetched()
        return item in self.related_objects

    def __iter__(self) -> "Iterator[MODEL]":
        self._raise_if_not_fetched()
        return self.related_objects.__iter__()

    def __len__(self) -> int:
        self._raise_if_not_fetched()
        return len(self.related_objects)

    def __bool__(self) -> bool:
        self._raise_if_not_fetched()
        return bool(self.related_objects)

    def __getitem__(self, item: int) -> MODEL:
        self._raise_if_not_fetched()
        return self.related_objects[item]

    def __await__(self) -> Generator[Any, None, List[MODEL]]:
        return self._query.__await__()

    async def __aiter__(self) -> AsyncGenerator[Any, MODEL]:
        if not self._fetched:
            self._set_result_for_query(await self)
        for val in self.related_objects:
            yield val


[docs]
    def filter(self, *args: "Q", **kwargs: Any) -> "QuerySet[MODEL]":
        """
        Returns a QuerySet with related elements filtered by args/kwargs.
        """
        return self._query.filter(*args, **kwargs)



[docs]
    def all(self) -> "QuerySet[MODEL]":
        """
        Returns a QuerySet with all related elements.
        """
        return self._query



[docs]
    def order_by(self, *orderings: str) -> "QuerySet[MODEL]":
        """
        Returns a QuerySet related elements in order.
        """
        return self._query.order_by(*orderings)



[docs]
    def limit(self, limit: int) -> "QuerySet[MODEL]":
        """
        Returns a QuerySet with at most «limit» related elements.
        """
        return self._query.limit(limit)



[docs]
    def offset(self, offset: int) -> "QuerySet[MODEL]":
        """
        Returns a QuerySet with all related elements offset by «offset».
        """
        return self._query.offset(offset)


    def _set_result_for_query(self, sequence: List[MODEL], attr: Optional[str] = None) -> None:
        self._fetched = True
        self.related_objects = sequence
        if attr:
            setattr(self.instance, attr, sequence)

    def _raise_if_not_fetched(self) -> None:
        if not self._fetched:
            raise NoValuesFetched(
                "No values were fetched for this relation, first use .fetch_related()"
            )




[docs]
class ManyToManyRelation(ReverseRelation[MODEL]):
    """
    Many-to-many relation container for :func:`.ManyToManyField`.
    """

    def __init__(self, instance: "Model", m2m_field: "ManyToManyFieldInstance[MODEL]") -> None:
        super().__init__(m2m_field.related_model, m2m_field.related_name, instance, "pk")
        self.field = m2m_field
        self.instance = instance


[docs]
    async def add(self, *instances: MODEL, using_db: "Optional[BaseDBAsyncClient]" = None) -> None:
        """
        Adds one or more of ``instances`` to the relation.

        If it is already added, it will be silently ignored.

        :raises OperationalError: If Object to add is not saved.
        """
        if not instances:
            return
        if not self.instance._saved_in_db:
            raise OperationalError(f"You should first call .save() on {self.instance}")
        db = using_db or self.remote_model._meta.db
        pk_formatting_func = type(self.instance)._meta.pk.to_db_value
        related_pk_formatting_func = type(instances[0])._meta.pk.to_db_value
        pk_b = pk_formatting_func(self.instance.pk, self.instance)
        pks_f: list = []
        for instance_to_add in instances:
            if not instance_to_add._saved_in_db:
                raise OperationalError(f"You should first call .save() on {instance_to_add}")
            pk_f = related_pk_formatting_func(instance_to_add.pk, instance_to_add)
            pks_f.append(pk_f)
        through_table = Table(self.field.through)
        backward_key, forward_key = self.field.backward_key, self.field.forward_key
        backward_field, forward_field = through_table[backward_key], through_table[forward_key]
        select_query = (
            db.query_class.from_(through_table).where(backward_field == pk_b).select(forward_key)
        )
        criterion = forward_field == pks_f[0] if len(pks_f) == 1 else forward_field.isin(pks_f)
        select_query = select_query.where(criterion)

        _, already_existing_relations_raw = await db.execute_query(str(select_query))
        already_existing_forward_pks = {
            related_pk_formatting_func(r[forward_key], self.instance)
            for r in already_existing_relations_raw
        }

        if pks_f_to_insert := set(pks_f) - already_existing_forward_pks:
            query = db.query_class.into(through_table).columns(forward_field, backward_field)
            for pk_f in pks_f_to_insert:
                query = query.insert(pk_f, pk_b)
            await db.execute_query(str(query))



[docs]
    async def clear(self, using_db: "Optional[BaseDBAsyncClient]" = None) -> None:
        """
        Clears ALL relations.
        """
        await self._remove_or_clear(using_db=using_db)



[docs]
    async def remove(
        self, *instances: MODEL, using_db: "Optional[BaseDBAsyncClient]" = None
    ) -> None:
        """
        Removes one or more of ``instances`` from the relation.

        :raises OperationalError: remove() was called with no instances.
        """
        if not instances:
            raise OperationalError("remove() called on no instances")
        await self._remove_or_clear(instances, using_db)


    async def _remove_or_clear(
        self,
        instances: Optional[Tuple[MODEL, ...]] = None,
        using_db: "Optional[BaseDBAsyncClient]" = None,
    ) -> None:
        db = using_db or self.remote_model._meta.db
        through_table = Table(self.field.through)
        pk_formatting_func = type(self.instance)._meta.pk.to_db_value

        condition = through_table[self.field.backward_key] == pk_formatting_func(
            self.instance.pk, self.instance
        )
        if instances:
            related_pk_formatting_func = type(instances[0])._meta.pk.to_db_value
            if len(instances) == 1:
                condition &= through_table[self.field.forward_key] == related_pk_formatting_func(
                    instances[0].pk, instances[0]
                )
            else:
                condition &= through_table[self.field.forward_key].isin(
                    [related_pk_formatting_func(i.pk, i) for i in instances]
                )
        query = db.query_class.from_(through_table).where(condition).delete()
        await db.execute_query(str(query))



class RelationalField(Field[MODEL]):
    has_db_field = False

    def __init__(
        self,
        related_model: "Type[MODEL]",
        to_field: Optional[str] = None,
        db_constraint: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.related_model: "Type[MODEL]" = related_model
        self.to_field: str = to_field  # type: ignore
        self.to_field_instance: Field = None  # type: ignore
        self.db_constraint = db_constraint

    if TYPE_CHECKING:

        @overload
        def __get__(self, instance: None, owner: Type["Model"]) -> "RelationalField[MODEL]": ...

        @overload
        def __get__(self, instance: "Model", owner: Type["Model"]) -> MODEL: ...

        def __get__(
            self, instance: Optional["Model"], owner: Type["Model"]
        ) -> "RelationalField[MODEL] | MODEL": ...

        def __set__(self, instance: "Model", value: MODEL) -> None: ...

    def describe(self, serializable: bool) -> dict:
        desc = super().describe(serializable)
        desc["db_constraint"] = self.db_constraint
        del desc["db_column"]
        return desc

    @classmethod
    def validate_model_name(cls, model_name: str) -> None:
        if len(model_name.split(".")) != 2:
            field_type = cls.__name__.replace("Instance", "")
            raise ConfigurationError(f'{field_type} accepts model name in format "app.Model"')


class ForeignKeyFieldInstance(RelationalField[MODEL]):
    def __init__(
        self,
        model_name: str,
        related_name: Union[Optional[str], Literal[False]] = None,
        on_delete: OnDelete = CASCADE,
        **kwargs: Any,
    ) -> None:
        super().__init__(None, **kwargs)  # type: ignore
        self.validate_model_name(model_name)
        self.model_name = model_name
        self.related_name = related_name
        if on_delete not in set(OnDelete):
            raise ConfigurationError(
                "on_delete can only be CASCADE, RESTRICT, SET_NULL, SET_DEFAULT or NO_ACTION"
            )
        if on_delete == SET_NULL and not bool(kwargs.get("null")):
            raise ConfigurationError("If on_delete is SET_NULL, then field must have null=True set")
        self.on_delete = on_delete

    def describe(self, serializable: bool) -> dict:
        desc = super().describe(serializable)
        desc["raw_field"] = self.source_field
        desc["on_delete"] = str(self.on_delete)
        return desc


class BackwardFKRelation(RelationalField[MODEL]):
    def __init__(
        self,
        field_type: "Type[MODEL]",
        relation_field: str,
        relation_source_field: str,
        null: bool,
        description: Optional[str],
        **kwargs: Any,
    ) -> None:
        super().__init__(field_type, null=null, **kwargs)
        self.relation_field: str = relation_field
        self.relation_source_field: str = relation_source_field
        self.description: Optional[str] = description


class OneToOneFieldInstance(ForeignKeyFieldInstance[MODEL]):
    def __init__(
        self,
        model_name: str,
        related_name: Union[Optional[str], Literal[False]] = None,
        on_delete: OnDelete = CASCADE,
        **kwargs: Any,
    ) -> None:
        self.validate_model_name(model_name)
        super().__init__(model_name, related_name, on_delete, unique=True, **kwargs)


class BackwardOneToOneRelation(BackwardFKRelation[MODEL]):
    pass


class ManyToManyFieldInstance(RelationalField[MODEL]):
    field_type = ManyToManyRelation

    def __init__(
        self,
        model_name: str,
        through: Optional[str] = None,
        forward_key: Optional[str] = None,
        backward_key: str = "",
        related_name: str = "",
        on_delete: OnDelete = CASCADE,
        field_type: "Type[MODEL]" = None,  # type: ignore
        create_unique_index: bool = True,
        **kwargs: Any,
    ) -> None:
        # TODO: rename through to through_table
        # TODO: add through to use a Model
        super().__init__(field_type, **kwargs)
        self.validate_model_name(model_name)
        self.model_name: str = model_name
        self.related_name: str = related_name
        self.forward_key: str = forward_key or f"{model_name.split('.')[1].lower()}_id"
        self.backward_key: str = backward_key
        self.through: str = through  # type: ignore
        self._generated: bool = False
        self.on_delete = on_delete
        self.create_unique_index = create_unique_index

    def describe(self, serializable: bool) -> dict:
        desc = super().describe(serializable)
        desc["model_name"] = self.model_name
        desc["related_name"] = self.related_name
        desc["forward_key"] = self.forward_key
        desc["backward_key"] = self.backward_key
        desc["through"] = self.through
        desc["on_delete"] = str(self.on_delete)
        desc["_generated"] = self._generated
        return desc


@overload
def OneToOneField(
    model_name: str,
    related_name: Union[Optional[str], Literal[False]] = None,
    on_delete: OnDelete = CASCADE,
    db_constraint: bool = True,
    *,
    null: Literal[True],
    **kwargs: Any,
) -> "OneToOneNullableRelation[MODEL]": ...


@overload
def OneToOneField(
    model_name: str,
    related_name: Union[Optional[str], Literal[False]] = None,
    on_delete: OnDelete = CASCADE,
    db_constraint: bool = True,
    null: Literal[False] = False,
    **kwargs: Any,
) -> "OneToOneRelation[MODEL]": ...



[docs]
def OneToOneField(
    model_name: str,
    related_name: Union[Optional[str], Literal[False]] = None,
    on_delete: OnDelete = CASCADE,
    db_constraint: bool = True,
    null: bool = False,
    **kwargs: Any,
) -> "OneToOneRelation[MODEL] | OneToOneNullableRelation[MODEL]":
    """
    OneToOne relation field.

    This field represents a foreign key relation to another model.

    See :ref:`one_to_one` for usage information.

    You must provide the following:

    ``model_name``:
        The name of the related model in a :samp:`'{app}.{model}'` format.

    The following is optional:

    ``related_name``:
        The attribute name on the related model to reverse resolve the foreign key.
    ``on_delete``:
        One of:
            ``field.CASCADE``:
                Indicate that the model should be cascade deleted if related model gets deleted.
            ``field.RESTRICT``:
                Indicate that the related model delete will be restricted as long as a
                foreign key points to it.
            ``field.SET_NULL``:
                Resets the field to NULL in case the related model gets deleted.
                Can only be set if field has ``null=True`` set.
            ``field.SET_DEFAULT``:
                Resets the field to ``default`` value in case the related model gets deleted.
                Can only be set is field has a ``default`` set.
            ``field.NO_ACTION``:
                Take no action.
    ``to_field``:
        The attribute name on the related model to establish foreign key relationship.
        If not set, pk is used
    ``db_constraint``:
        Controls whether or not a constraint should be created in the database for this foreign key.
        The default is True, and that’s almost certainly what you want; setting this to False can be very bad for data integrity.
    """

    return OneToOneFieldInstance(
        model_name, related_name, on_delete, db_constraint=db_constraint, null=null, **kwargs
    )



@overload
def ForeignKeyField(
    model_name: str,
    related_name: Union[Optional[str], Literal[False]] = None,
    on_delete: OnDelete = CASCADE,
    db_constraint: bool = True,
    *,
    null: Literal[True],
    **kwargs: Any,
) -> "ForeignKeyNullableRelation[MODEL]": ...


@overload
def ForeignKeyField(
    model_name: str,
    related_name: Union[Optional[str], Literal[False]] = None,
    on_delete: OnDelete = CASCADE,
    db_constraint: bool = True,
    null: Literal[False] = False,
    **kwargs: Any,
) -> "ForeignKeyRelation[MODEL]": ...



[docs]
def ForeignKeyField(
    model_name: str,
    related_name: Union[Optional[str], Literal[False]] = None,
    on_delete: OnDelete = CASCADE,
    db_constraint: bool = True,
    null: bool = False,
    **kwargs: Any,
) -> "ForeignKeyRelation[MODEL] | ForeignKeyNullableRelation[MODEL]":
    """
    ForeignKey relation field.

    This field represents a foreign key relation to another model.

    See :ref:`foreign_key` for usage information.

    You must provide the following:

    ``model_name``:
        The name of the related model in a :samp:`'{app}.{model}'` format.

    The following is optional:

    ``related_name``:
        The attribute name on the related model to reverse resolve the foreign key.
    ``on_delete``:
        One of:
            ``field.CASCADE``:
                Indicate that the model should be cascade deleted if related model gets deleted.
            ``field.RESTRICT``:
                Indicate that the related model delete will be restricted as long as a
                foreign key points to it.
            ``field.SET_NULL``:
                Resets the field to NULL in case the related model gets deleted.
                Can only be set if field has ``null=True`` set.
            ``field.SET_DEFAULT``:
                Resets the field to ``default`` value in case the related model gets deleted.
                Can only be set is field has a ``default`` set.
            ``field.NO_ACTION``:
                Take no action.
    ``to_field``:
        The attribute name on the related model to establish foreign key relationship.
        If not set, pk is used
    ``db_constraint``:
        Controls whether or not a constraint should be created in the database for this foreign key.
        The default is True, and that’s almost certainly what you want; setting this to False can be very bad for data integrity.
    """

    return ForeignKeyFieldInstance(
        model_name, related_name, on_delete, db_constraint=db_constraint, null=null, **kwargs
    )




[docs]
def ManyToManyField(
    model_name: str,
    through: Optional[str] = None,
    forward_key: Optional[str] = None,
    backward_key: str = "",
    related_name: str = "",
    on_delete: OnDelete = CASCADE,
    db_constraint: bool = True,
    create_unique_index: bool = True,
    **kwargs: Any,
) -> "ManyToManyRelation[Any]":
    """
    ManyToMany relation field.

    This field represents a many-to-many between this model and another model.

    See :ref:`many_to_many` for usage information.

    You must provide the following:

    ``model_name``:
        The name of the related model in a :samp:`'{app}.{model}'` format.

    The following is optional:

    ``through``:
        The DB table that represents the through table.
        The default is normally safe.
    ``forward_key``:
        The forward lookup key on the through table.
        The default is normally safe.
    ``backward_key``:
        The backward lookup key on the through table.
        The default is normally safe.
    ``related_name``:
        The attribute name on the related model to reverse resolve the many to many.
    ``db_constraint``:
        Controls whether or not a constraint should be created in the database for this foreign key.
        The default is True, and that’s almost certainly what you want; setting this to False can be very bad for data integrity.
    ``on_delete``:
        One of:
            ``field.CASCADE``:
                Indicate that the model should be cascade deleted if related model gets deleted.
            ``field.RESTRICT``:
                Indicate that the related model delete will be restricted as long as a
                foreign key points to it.
            ``field.SET_NULL``:
                Resets the field to NULL in case the related model gets deleted.
                Can only be set if field has ``null=True`` set.
            ``field.SET_DEFAULT``:
                Resets the field to ``default`` value in case the related model gets deleted.
                Can only be set is field has a ``default`` set.
            ``field.NO_ACTION``:
                Take no action.
    ``create_unique_index``:
        Controls whether or not a unique index should be created in the database to speed up select queries.
        The default is True. If you want to allow repeat records, set this to False.
    """

    return ManyToManyFieldInstance(  # type: ignore
        model_name,
        through,
        forward_key,
        backward_key,
        related_name,
        on_delete=on_delete,
        db_constraint=db_constraint,
        create_unique_index=create_unique_index,
        **kwargs,
    )



OneToOneNullableRelation = Optional[OneToOneFieldInstance[MODEL]]
"""
Type hint for the result of accessing the :func:`.OneToOneField` field in the model
when obtained model can be nullable.
"""

OneToOneRelation = OneToOneFieldInstance[MODEL]
"""
Type hint for the result of accessing the :func:`.OneToOneField` field in the model.
"""

ForeignKeyNullableRelation = Optional[ForeignKeyFieldInstance[MODEL]]
"""
Type hint for the result of accessing the :func:`.ForeignKeyField` field in the model
when obtained model can be nullable.
"""

ForeignKeyRelation = ForeignKeyFieldInstance[MODEL]
"""
Type hint for the result of accessing the :func:`.ForeignKeyField` field in the model.
"""