122 lines
4.6 KiB
Python
122 lines
4.6 KiB
Python
from typing import Callable, Dict, List, TYPE_CHECKING, Tuple, Type, cast
|
|
|
|
from ormar.models.mixins.relation_mixin import RelationMixin
|
|
|
|
if TYPE_CHECKING: # pragma: no cover
|
|
from ormar.fields import ForeignKeyField, ManyToManyField
|
|
|
|
|
|
class PrefetchQueryMixin(RelationMixin):
|
|
"""
|
|
Used in PrefetchQuery to extract ids and names of models to prefetch.
|
|
"""
|
|
|
|
if TYPE_CHECKING: # pragma no cover
|
|
from ormar import Model
|
|
|
|
get_name: Callable # defined in NewBaseModel
|
|
|
|
@staticmethod
|
|
def get_clause_target_and_filter_column_name(
|
|
parent_model: Type["Model"],
|
|
target_model: Type["Model"],
|
|
reverse: bool,
|
|
related: str,
|
|
) -> Tuple[Type["Model"], str]:
|
|
"""
|
|
Returns Model on which query clause should be performed and name of the column.
|
|
|
|
:param parent_model: related model that the relation lead to
|
|
:type parent_model: Type[Model]
|
|
:param target_model: model on which query should be perfomed
|
|
:type target_model: Type[Model]
|
|
:param reverse: flag if the relation is reverse
|
|
:type reverse: bool
|
|
:param related: name of the relation field
|
|
:type related: str
|
|
:return: Model on which query clause should be performed and name of the column
|
|
:rtype: Tuple[Type[Model], str]
|
|
"""
|
|
if reverse:
|
|
field_name = parent_model.Meta.model_fields[related].get_related_name()
|
|
field = target_model.Meta.model_fields[field_name]
|
|
if field.is_multi:
|
|
field = cast(Type["ManyToManyField"], field)
|
|
field_name = field.default_target_field_name()
|
|
sub_field = field.through.Meta.model_fields[field_name]
|
|
return field.through, sub_field.get_alias()
|
|
return target_model, field.get_alias()
|
|
target_field = target_model.get_column_alias(target_model.Meta.pkname)
|
|
return target_model, target_field
|
|
|
|
@staticmethod
|
|
def get_column_name_for_id_extraction(
|
|
parent_model: Type["Model"], reverse: bool, related: str, use_raw: bool,
|
|
) -> str:
|
|
"""
|
|
Returns name of the column that should be used to extract ids from model.
|
|
Depending on the relation side it's either primary key column of parent model
|
|
or field name specified by related parameter.
|
|
|
|
:param parent_model: model from which id column should be extracted
|
|
:type parent_model: Type[Model]
|
|
:param reverse: flag if the relation is reverse
|
|
:type reverse: bool
|
|
:param related: name of the relation field
|
|
:type related: str
|
|
:param use_raw: flag if aliases or field names should be used
|
|
:type use_raw: bool
|
|
:return:
|
|
:rtype:
|
|
"""
|
|
if reverse:
|
|
column_name = parent_model.Meta.pkname
|
|
return (
|
|
parent_model.get_column_alias(column_name) if use_raw else column_name
|
|
)
|
|
column = parent_model.Meta.model_fields[related]
|
|
return column.get_alias() if use_raw else column.name
|
|
|
|
@classmethod
|
|
def get_related_field_name(cls, target_field: Type["ForeignKeyField"]) -> str:
|
|
"""
|
|
Returns name of the relation field that should be used in prefetch query.
|
|
This field is later used to register relation in prefetch query,
|
|
populate relations dict, and populate nested model in prefetch query.
|
|
|
|
:param target_field: relation field that should be used in prefetch
|
|
:type target_field: Type[BaseField]
|
|
:return: name of the field
|
|
:rtype: str
|
|
"""
|
|
if target_field.is_multi:
|
|
return cls.get_name()
|
|
if target_field.virtual:
|
|
return target_field.get_related_name()
|
|
return target_field.to.Meta.pkname
|
|
|
|
@classmethod
|
|
def get_filtered_names_to_extract(cls, prefetch_dict: Dict) -> List:
|
|
"""
|
|
Returns list of related fields names that should be followed to prefetch related
|
|
models from.
|
|
|
|
List of models is translated into dict to assure each model is extracted only
|
|
once in one query, that's why this function accepts prefetch_dict not list.
|
|
|
|
Only relations from current model are returned.
|
|
|
|
:param prefetch_dict: dictionary of fields to extract
|
|
:type prefetch_dict: Dict
|
|
:return: list of fields names to extract
|
|
:rtype: List
|
|
"""
|
|
related_to_extract = []
|
|
if prefetch_dict and prefetch_dict is not Ellipsis:
|
|
related_to_extract = [
|
|
related
|
|
for related in cls.extract_related_names()
|
|
if related in prefetch_dict
|
|
]
|
|
return related_to_extract
|