diff --git a/ormar/models/excludable.py b/ormar/models/excludable.py index b754eb4..b8832d2 100644 --- a/ormar/models/excludable.py +++ b/ormar/models/excludable.py @@ -12,6 +12,20 @@ class Excludable: include: Set = field(default_factory=set) exclude: Set = field(default_factory=set) + @property + def include_all(self): + return ... in self.include + + @property + def exclude_all(self): + return ... in self.exclude + + def get_copy(self) -> "Excludable": + _copy = self.__class__() + _copy.include = {x for x in self.include} + _copy.exclude = {x for x in self.exclude} + return _copy + def set_values(self, value: Set, is_exclude: bool) -> None: prop = "exclude" if is_exclude else "include" if ... in getattr(self, prop) or ... in value: @@ -38,15 +52,22 @@ class ExcludableItems: def __init__(self) -> None: self.items: Dict[str, Excludable] = dict() + @classmethod + def from_excludable(cls, other: "ExcludableItems") -> "ExcludableItems": + new_excludable = cls() + for key, value in other.items.items(): + new_excludable.items[key] = value.get_copy() + return new_excludable + def get(self, model_cls: Type["Model"], alias: str = "") -> Excludable: key = f"{alias + '_' if alias else ''}{model_cls.get_name(lower=True)}" return self.items.get(key, Excludable()) def build( - self, - items: Union[List[str], str, Tuple[str], Set[str], Dict], - model_cls: Type["Model"], - is_exclude: bool = False, + self, + items: Union[List[str], str, Tuple[str], Set[str], Dict], + model_cls: Type["Model"], + is_exclude: bool = False, ) -> None: if isinstance(items, str): @@ -75,7 +96,7 @@ class ExcludableItems: ) def _set_excludes( - self, items: Set, model_name: str, is_exclude: bool, alias: str = "" + self, items: Set, model_name: str, is_exclude: bool, alias: str = "" ) -> None: key = f"{alias + '_' if alias else ''}{model_name}" @@ -86,13 +107,13 @@ class ExcludableItems: self.items[key] = excludable def _traverse_dict( # noqa: CFQ002 - self, - values: Dict, - source_model: Type["Model"], - model_cls: Type["Model"], - is_exclude: bool, - related_items: List = None, - alias: str = "", + self, + values: Dict, + source_model: Type["Model"], + model_cls: Type["Model"], + is_exclude: bool, + related_items: List = None, + alias: str = "", ) -> None: self_fields = set() @@ -144,7 +165,7 @@ class ExcludableItems: ) def _traverse_list( - self, values: Set[str], model_cls: Type["Model"], is_exclude: bool + self, values: Set[str], model_cls: Type["Model"], is_exclude: bool ) -> None: # here we have only nested related keys diff --git a/ormar/models/mixins/excludable_mixin.py b/ormar/models/mixins/excludable_mixin.py index 4b096d9..4b25035 100644 --- a/ormar/models/mixins/excludable_mixin.py +++ b/ormar/models/mixins/excludable_mixin.py @@ -9,9 +9,10 @@ from typing import ( TYPE_CHECKING, Type, TypeVar, - Union, + Union, cast, ) +from ormar.models.excludable import ExcludableItems from ormar.models.mixins.relation_mixin import RelationMixin from ormar.queryset.utils import translate_list_to_dict, update @@ -35,7 +36,7 @@ class ExcludableMixin(RelationMixin): @staticmethod def get_child( - items: Union[Set, Dict, None], key: str = None + items: Union[Set, Dict, None], key: str = None ) -> Union[Set, Dict, None]: """ Used to get nested dictionaries keys if they exists otherwise returns @@ -53,7 +54,7 @@ class ExcludableMixin(RelationMixin): @staticmethod def get_excluded( - exclude: Union[Set, Dict, None], key: str = None + exclude: Union[Set, Dict, None], key: str = None ) -> Union[Set, Dict, None]: """ Proxy to ExcludableMixin.get_child for exclusions. @@ -69,7 +70,7 @@ class ExcludableMixin(RelationMixin): @staticmethod def get_included( - include: Union[Set, Dict, None], key: str = None + include: Union[Set, Dict, None], key: str = None ) -> Union[Set, Dict, None]: """ Proxy to ExcludableMixin.get_child for inclusions. @@ -131,9 +132,9 @@ class ExcludableMixin(RelationMixin): @staticmethod def _populate_pk_column( - model: Union[Type["Model"], Type["ModelRow"]], - columns: List[str], - use_alias: bool = False, + model: Union[Type["Model"], Type["ModelRow"]], + columns: List[str], + use_alias: bool = False, ) -> List[str]: """ Adds primary key column/alias (depends on use_alias flag) to list of @@ -159,12 +160,13 @@ class ExcludableMixin(RelationMixin): @classmethod def own_table_columns( - cls, - model: Union[Type["Model"], Type["ModelRow"]], - fields: Optional[Union[Set, Dict]], - exclude_fields: Optional[Union[Set, Dict]], - use_alias: bool = False, + cls, + model: Union[Type["Model"], Type["ModelRow"]], + excludable: ExcludableItems, + alias: str = '', + use_alias: bool = False, ) -> List[str]: + # TODO update docstring """ Returns list of aliases or field names for given model. Aliases/names switch is use_alias flag. @@ -176,15 +178,12 @@ class ExcludableMixin(RelationMixin): :param model: model on columns are selected :type model: Type["Model"] - :param fields: set/dict of fields to include - :type fields: Optional[Union[Set, Dict]] - :param exclude_fields: set/dict of fields to exclude - :type exclude_fields: Optional[Union[Set, Dict]] :param use_alias: flag if aliases or field names should be used :type use_alias: bool :return: list of column field names or aliases :rtype: List[str] """ + model_excludable = excludable.get(model_cls=model, alias=alias) columns = [ model.get_column_name_from_alias(col.name) if not use_alias else col.name for col in model.Meta.table.columns @@ -193,17 +192,17 @@ class ExcludableMixin(RelationMixin): model.get_column_name_from_alias(col.name) for col in model.Meta.table.columns ] - if fields: + if model_excludable.include: columns = [ col for col, name in zip(columns, field_names) - if model.is_included(fields, name) + if model_excludable.is_included(name) ] - if exclude_fields: + if model_excludable.exclude: columns = [ col for col, name in zip(columns, field_names) - if not model.is_excluded(exclude_fields, name) + if not model_excludable.is_excluded(name) ] # always has to return pk column for ormar to work @@ -215,9 +214,9 @@ class ExcludableMixin(RelationMixin): @classmethod def _update_excluded_with_related_not_required( - cls, - exclude: Union["AbstractSetIntStr", "MappingIntStrAny", None], - nested: bool = False, + cls, + exclude: Union["AbstractSetIntStr", "MappingIntStrAny", None], + nested: bool = False, ) -> Union[Set, Dict]: """ Used during generation of the dict(). @@ -245,9 +244,9 @@ class ExcludableMixin(RelationMixin): @classmethod def get_names_to_exclude( - cls, - fields: Optional[Union[Dict, Set]] = None, - exclude_fields: Optional[Union[Dict, Set]] = None, + cls, + excludable: ExcludableItems, + alias: str ) -> Set: """ Returns a set of models field names that should be explicitly excluded @@ -259,33 +258,27 @@ class ExcludableMixin(RelationMixin): Used in parsing data from database rows that construct Models by initializing them with dicts constructed from those db rows. - :param fields: set/dict of fields to include - :type fields: Optional[Union[Set, Dict]] - :param exclude_fields: set/dict of fields to exclude - :type exclude_fields: Optional[Union[Set, Dict]] + :param alias: alias of current relation + :type alias: str + :param excludable: structure of fields to include and exclude + :type excludable: ExcludableItems :return: set of field names that should be excluded :rtype: Set """ + model = cast(Type["Model"], cls) + model_excludable = excludable.get(model_cls=model, alias=alias) fields_names = cls.extract_db_own_fields() - if fields and fields is not Ellipsis: - fields_to_keep = {name for name in fields if name in fields_names} + if model_excludable.include and model_excludable.include_all: + fields_to_keep = model_excludable.include.intersection(fields_names) else: fields_to_keep = fields_names fields_to_exclude = fields_names - fields_to_keep - if isinstance(exclude_fields, Set): + if model_excludable.exclude: fields_to_exclude = fields_to_exclude.union( - {name for name in exclude_fields if name in fields_names} + model_excludable.exclude.intersection(fields_names) ) - elif isinstance(exclude_fields, Dict): - new_to_exclude = { - name - for name in exclude_fields - if name in fields_names and exclude_fields[name] is Ellipsis - } - fields_to_exclude = fields_to_exclude.union(new_to_exclude) - fields_to_exclude = fields_to_exclude - {cls.Meta.pkname} return fields_to_exclude diff --git a/ormar/models/model_row.py b/ormar/models/model_row.py index 476e274..6a6cb0e 100644 --- a/ormar/models/model_row.py +++ b/ormar/models/model_row.py @@ -14,6 +14,7 @@ from typing import ( import sqlalchemy from ormar.models import NewBaseModel # noqa: I202 +from ormar.models.excludable import ExcludableItems from ormar.models.helpers.models import group_related_list @@ -33,8 +34,7 @@ class ModelRow(NewBaseModel): select_related: List = None, related_models: Any = None, related_field: Type["ForeignKeyField"] = None, - fields: Optional[Union[Dict, Set]] = None, - exclude_fields: Optional[Union[Dict, Set]] = None, + excludable: ExcludableItems = None, current_relation_str: str = "", ) -> Optional[T]: """ @@ -50,6 +50,8 @@ class ModelRow(NewBaseModel): where rows are populated in a different way as they do not have nested models in result. + :param excludable: structure of fields to include and exclude + :type excludable: ExcludableItems :param current_relation_str: name of the relation field :type current_relation_str: str :param source_model: model on which relation was defined @@ -62,12 +64,6 @@ class ModelRow(NewBaseModel): :type related_models: Union[List, Dict] :param related_field: field with relation declaration :type related_field: Type[ForeignKeyField] - :param fields: fields and related model fields to include - if provided only those are included - :type fields: Optional[Union[Dict, Set]] - :param exclude_fields: fields and related model fields to exclude - excludes the fields even if they are provided in fields - :type exclude_fields: Optional[Union[Dict, Set]] :return: returns model if model is populated from database :rtype: Optional[Model] """ @@ -75,6 +71,7 @@ class ModelRow(NewBaseModel): select_related = select_related or [] related_models = related_models or [] table_prefix = "" + excludable = excludable or ExcludableItems() if select_related: source_model = cast(Type[T], cls) @@ -87,12 +84,11 @@ class ModelRow(NewBaseModel): relation_field=related_field, ) - item = cls.populate_nested_models_from_row( + item = cls._populate_nested_models_from_row( item=item, row=row, related_models=related_models, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, current_relation_str=current_relation_str, source_model=source_model, ) @@ -100,28 +96,26 @@ class ModelRow(NewBaseModel): item=item, row=row, table_prefix=table_prefix, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable ) instance: Optional[T] = None if item.get(cls.Meta.pkname, None) is not None: item["__excluded__"] = cls.get_names_to_exclude( - fields=fields, exclude_fields=exclude_fields + excludable=excludable, alias=table_prefix ) instance = cast(T, cls(**item)) instance.set_save_status(True) return instance @classmethod - def populate_nested_models_from_row( # noqa: CFQ002 + def _populate_nested_models_from_row( # noqa: CFQ002 cls, item: dict, row: sqlalchemy.engine.ResultProxy, source_model: Type[T], related_models: Any, - fields: Optional[Union[Dict, Set]] = None, - exclude_fields: Optional[Union[Dict, Set]] = None, + excludable: ExcludableItems, current_relation_str: str = None, ) -> dict: """ @@ -134,6 +128,8 @@ class ModelRow(NewBaseModel): Recurrently calls from_row method on nested instances and create nested instances. In the end those instances are added to the final model dictionary. + :param excludable: structure of fields to include and exclude + :type excludable: ExcludableItems :param source_model: source model from which relation started :type source_model: Type[Model] :param current_relation_str: joined related parts into one string @@ -144,12 +140,6 @@ class ModelRow(NewBaseModel): :type row: sqlalchemy.engine.result.ResultProxy :param related_models: list or dict of related models :type related_models: Union[Dict, List] - :param fields: fields and related model fields to include - - if provided only those are included - :type fields: Optional[Union[Dict, Set]] - :param exclude_fields: fields and related model fields to exclude - excludes the fields even if they are provided in fields - :type exclude_fields: Optional[Union[Dict, Set]] :return: dictionary with keys corresponding to model fields names and values are database values :rtype: Dict @@ -163,8 +153,6 @@ class ModelRow(NewBaseModel): ) field = cls.Meta.model_fields[related] field = cast(Type["ForeignKeyField"], field) - fields = cls.get_included(fields, related) - exclude_fields = cls.get_excluded(exclude_fields, related) model_cls = field.to remainder = None @@ -174,8 +162,7 @@ class ModelRow(NewBaseModel): row, related_models=remainder, related_field=field, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, current_relation_str=relation_str, source_model=source_model, ) @@ -188,8 +175,7 @@ class ModelRow(NewBaseModel): row=row, related=related, through_name=through_name, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable ) item[through_name] = through_child setattr(child, through_name, through_child) @@ -203,35 +189,29 @@ class ModelRow(NewBaseModel): row: sqlalchemy.engine.ResultProxy, through_name: str, related: str, - fields: Optional[Union[Dict, Set]] = None, - exclude_fields: Optional[Union[Dict, Set]] = None, + excludable: ExcludableItems ) -> Dict: - # TODO: fix excludes and includes - fields = cls.get_included(fields, through_name) - # exclude_fields = cls.get_excluded(exclude_fields, through_name) + # TODO: fix excludes and includes and docstring model_cls = cls.Meta.model_fields[through_name].to - exclude_fields = model_cls.extract_related_names() table_prefix = cls.Meta.alias_manager.resolve_relation_alias( from_model=cls, relation_name=related ) child = model_cls.extract_prefixed_table_columns( item={}, row=row, - table_prefix=table_prefix, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, + table_prefix=table_prefix ) return child @classmethod - def extract_prefixed_table_columns( # noqa CCR001 + def extract_prefixed_table_columns( cls, item: dict, row: sqlalchemy.engine.result.ResultProxy, table_prefix: str, - fields: Optional[Union[Dict, Set]] = None, - exclude_fields: Optional[Union[Dict, Set]] = None, - ) -> dict: + excludable: ExcludableItems + ) -> Dict: """ Extracts own fields from raw sql result, using a given prefix. Prefix changes depending on the table's position in a join. @@ -244,6 +224,8 @@ class ModelRow(NewBaseModel): Used in Model.from_row and PrefetchQuery._populate_rows methods. + :param excludable: structure of fields to include and exclude + :type excludable: ExcludableItems :param item: dictionary of already populated nested models, otherwise empty dict :type item: Dict :param row: raw result row from the database @@ -252,12 +234,6 @@ class ModelRow(NewBaseModel): each pair of tables have own prefix (two of them depending on direction) - used in joins to allow multiple joins to the same table. :type table_prefix: str - :param fields: fields and related model fields to include - - if provided only those are included - :type fields: Optional[Union[Dict, Set]] - :param exclude_fields: fields and related model fields to exclude - excludes the fields even if they are provided in fields - :type exclude_fields: Optional[Union[Dict, Set]] :return: dictionary with keys corresponding to model fields names and values are database values :rtype: Dict @@ -267,8 +243,8 @@ class ModelRow(NewBaseModel): selected_columns = cls.own_table_columns( model=cls, - fields=fields or {}, - exclude_fields=exclude_fields or {}, + excludable=excludable, + alias=table_prefix, use_alias=False, ) diff --git a/ormar/queryset/join.py b/ormar/queryset/join.py index f18c81d..a6f1e93 100644 --- a/ormar/queryset/join.py +++ b/ormar/queryset/join.py @@ -16,6 +16,7 @@ from sqlalchemy import text import ormar # noqa I100 from ormar.exceptions import RelationshipInstanceError +from ormar.models.excludable import ExcludableItems from ormar.relations import AliasManager if TYPE_CHECKING: # pragma no cover @@ -29,8 +30,7 @@ class SqlJoin: used_aliases: List, select_from: sqlalchemy.sql.select, columns: List[sqlalchemy.Column], - fields: Optional[Union[Set, Dict]], - exclude_fields: Optional[Union[Set, Dict]], + excludable: ExcludableItems, order_columns: Optional[List["OrderAction"]], sorted_orders: OrderedDict, main_model: Type["Model"], @@ -44,8 +44,7 @@ class SqlJoin: self.related_models = related_models or [] self.select_from = select_from self.columns = columns - self.fields = fields - self.exclude_fields = exclude_fields + self.excludable=excludable self.order_columns = order_columns self.sorted_orders = sorted_orders self.main_model = main_model @@ -200,10 +199,7 @@ class SqlJoin: used_aliases=self.used_aliases, select_from=self.select_from, columns=self.columns, - fields=self.main_model.get_excluded(self.fields, related_name), - exclude_fields=self.main_model.get_excluded( - self.exclude_fields, related_name - ), + excludable=self.excludable, order_columns=self.order_columns, sorted_orders=self.sorted_orders, main_model=self.next_model, @@ -303,8 +299,8 @@ class SqlJoin: # TODO: fix fields and exclusions for through model? self_related_fields = self.next_model.own_table_columns( model=self.next_model, - fields=self.fields, - exclude_fields=self.exclude_fields, + excludable=self.excludable, + alias=self.next_alias, use_alias=True, ) self.columns.extend( diff --git a/ormar/queryset/prefetch_query.py b/ormar/queryset/prefetch_query.py index 533f92c..88fc8e3 100644 --- a/ormar/queryset/prefetch_query.py +++ b/ormar/queryset/prefetch_query.py @@ -13,6 +13,7 @@ from typing import ( ) import ormar +from ormar.models.excludable import ExcludableItems from ormar.queryset.clause import QueryClause from ormar.queryset.query import Query from ormar.queryset.utils import extract_models_to_dict_of_lists, translate_list_to_dict @@ -24,7 +25,7 @@ if TYPE_CHECKING: # pragma: no cover def add_relation_field_to_fields( - fields: Union[Set[Any], Dict[Any, Any], None], related_field_name: str + fields: Union[Set[Any], Dict[Any, Any], None], related_field_name: str ) -> Union[Set[Any], Dict[Any, Any], None]: """ Adds related field into fields to include as otherwise it would be skipped. @@ -73,12 +74,12 @@ def sort_models(models: List["Model"], orders_by: Dict) -> List["Model"]: def set_children_on_model( # noqa: CCR001 - model: "Model", - related: str, - children: Dict, - model_id: int, - models: Dict, - orders_by: Dict, + model: "Model", + related: str, + children: Dict, + model_id: int, + models: Dict, + orders_by: Dict, ) -> None: """ Extract ids of child models by given relation id key value. @@ -123,21 +124,19 @@ class PrefetchQuery: """ def __init__( # noqa: CFQ002 - self, - model_cls: Type["Model"], - fields: Optional[Union[Dict, Set]], - exclude_fields: Optional[Union[Dict, Set]], - prefetch_related: List, - select_related: List, - orders_by: List["OrderAction"], + self, + model_cls: Type["Model"], + excludable: ExcludableItems, + prefetch_related: List, + select_related: List, + orders_by: List["OrderAction"], ) -> None: self.model = model_cls self.database = self.model.Meta.database self._prefetch_related = prefetch_related self._select_related = select_related - self._exclude_columns = exclude_fields - self._columns = fields + self.excludable = excludable self.already_extracted: Dict = dict() self.models: Dict = {} self.select_dict = translate_list_to_dict(self._select_related) @@ -148,7 +147,7 @@ class PrefetchQuery: ) async def prefetch_related( - self, models: Sequence["Model"], rows: List + self, models: Sequence["Model"], rows: List ) -> Sequence["Model"]: """ Main entry point for prefetch_query. @@ -173,7 +172,7 @@ class PrefetchQuery: return await self._prefetch_related_models(models=models, rows=rows) def _extract_ids_from_raw_data( - self, parent_model: Type["Model"], column_name: str + self, parent_model: Type["Model"], column_name: str ) -> Set: """ Iterates over raw rows and extract id values of relation columns by using @@ -196,7 +195,7 @@ class PrefetchQuery: return list_of_ids def _extract_ids_from_preloaded_models( - self, parent_model: Type["Model"], column_name: str + self, parent_model: Type["Model"], column_name: str ) -> Set: """ Extracts relation ids from already populated models if they were included @@ -219,7 +218,7 @@ class PrefetchQuery: return list_of_ids def _extract_required_ids( - self, parent_model: Type["Model"], reverse: bool, related: str, + self, parent_model: Type["Model"], reverse: bool, related: str, ) -> Set: """ Delegates extraction of the fields to either get ids from raw sql response @@ -253,11 +252,11 @@ class PrefetchQuery: ) def _get_filter_for_prefetch( - self, - parent_model: Type["Model"], - target_model: Type["Model"], - reverse: bool, - related: str, + self, + parent_model: Type["Model"], + target_model: Type["Model"], + reverse: bool, + related: str, ) -> List: """ Populates where clause with condition to return only models within the @@ -298,7 +297,7 @@ class PrefetchQuery: return [] def _populate_nested_related( - self, model: "Model", prefetch_dict: Dict, orders_by: Dict, + self, model: "Model", prefetch_dict: Dict, orders_by: Dict, ) -> "Model": """ Populates all related models children of parent model that are @@ -342,7 +341,7 @@ class PrefetchQuery: return model async def _prefetch_related_models( - self, models: Sequence["Model"], rows: List + self, models: Sequence["Model"], rows: List ) -> Sequence["Model"]: """ Main method of the query. @@ -366,8 +365,6 @@ class PrefetchQuery: select_dict = translate_list_to_dict(self._select_related) prefetch_dict = translate_list_to_dict(self._prefetch_related) target_model = self.model - fields = self._columns - exclude_fields = self._exclude_columns orders_by = self.order_dict for related in prefetch_dict.keys(): await self._extract_related_models( @@ -375,8 +372,7 @@ class PrefetchQuery: target_model=target_model, prefetch_dict=prefetch_dict.get(related, {}), select_dict=select_dict.get(related, {}), - fields=fields, - exclude_fields=exclude_fields, + excludable=self.excludable, orders_by=orders_by.get(related, {}), ) final_models = [] @@ -389,14 +385,13 @@ class PrefetchQuery: return models async def _extract_related_models( # noqa: CFQ002, CCR001 - self, - related: str, - target_model: Type["Model"], - prefetch_dict: Dict, - select_dict: Dict, - fields: Union[Set[Any], Dict[Any, Any], None], - exclude_fields: Union[Set[Any], Dict[Any, Any], None], - orders_by: Dict, + self, + related: str, + target_model: Type["Model"], + prefetch_dict: Dict, + select_dict: Dict, + excludable: ExcludableItems, + orders_by: Dict, ) -> None: """ Constructs queries with required ids and extracts data with fields that should @@ -424,8 +419,6 @@ class PrefetchQuery: :return: None :rtype: None """ - fields = target_model.get_included(fields, related) - exclude_fields = target_model.get_excluded(exclude_fields, related) target_field = target_model.Meta.model_fields[related] target_field = cast(Type["ForeignKeyField"], target_field) reverse = False @@ -450,14 +443,11 @@ class PrefetchQuery: related_field_name = parent_model.get_related_field_name( target_field=target_field ) - fields = add_relation_field_to_fields( - fields=fields, related_field_name=related_field_name - ) table_prefix, rows = await self._run_prefetch_query( target_field=target_field, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, filter_clauses=filter_clauses, + related_field_name=related_field_name ) else: rows = [] @@ -472,8 +462,7 @@ class PrefetchQuery: select_dict=self._get_select_related_if_apply( subrelated, select_dict ), - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, orders_by=self._get_select_related_if_apply(subrelated, orders_by), ) @@ -483,8 +472,7 @@ class PrefetchQuery: parent_model=parent_model, target_field=target_field, table_prefix=table_prefix, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, prefetch_dict=prefetch_dict, orders_by=orders_by, ) @@ -496,11 +484,11 @@ class PrefetchQuery: ) async def _run_prefetch_query( - self, - target_field: Type["BaseField"], - fields: Union[Set[Any], Dict[Any, Any], None], - exclude_fields: Union[Set[Any], Dict[Any, Any], None], - filter_clauses: List, + self, + target_field: Type["BaseField"], + excludable: ExcludableItems, + filter_clauses: List, + related_field_name: str ) -> Tuple[str, List]: """ Actually runs the queries against the database and populates the raw response @@ -511,10 +499,6 @@ class PrefetchQuery: :param target_field: ormar field with relation definition :type target_field: Type["BaseField"] - :param fields: fields to include - :type fields: Union[Set[Any], Dict[Any, Any], None] - :param exclude_fields: fields to exclude - :type exclude_fields: Union[Set[Any], Dict[Any, Any], None] :param filter_clauses: list of clauses, actually one clause with ids of relation :type filter_clauses: List[sqlalchemy.sql.elements.TextClause] :return: table prefix and raw rows from sql response @@ -533,6 +517,11 @@ class PrefetchQuery: ) self.already_extracted.setdefault(target_name, {})["prefix"] = table_prefix + model_excludable = excludable.get(model_cls=target_model, alias=table_prefix) + if model_excludable.include and not model_excludable.is_included( + related_field_name): + model_excludable.set_values({related_field_name}, is_exclude=False) + qry = Query( model_cls=query_target, select_related=select_related, @@ -540,8 +529,7 @@ class PrefetchQuery: exclude_clauses=[], offset=None, limit_count=None, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, order_bys=None, limit_raw_sql=False, ) @@ -571,7 +559,7 @@ class PrefetchQuery: ) def _update_already_loaded_rows( # noqa: CFQ002 - self, target_field: Type["BaseField"], prefetch_dict: Dict, orders_by: Dict, + self, target_field: Type["BaseField"], prefetch_dict: Dict, orders_by: Dict, ) -> None: """ Updates models that are already loaded, usually children of children. @@ -590,15 +578,14 @@ class PrefetchQuery: ) def _populate_rows( # noqa: CFQ002 - self, - rows: List, - target_field: Type["ForeignKeyField"], - parent_model: Type["Model"], - table_prefix: str, - fields: Union[Set[Any], Dict[Any, Any], None], - exclude_fields: Union[Set[Any], Dict[Any, Any], None], - prefetch_dict: Dict, - orders_by: Dict, + self, + rows: List, + target_field: Type["ForeignKeyField"], + parent_model: Type["Model"], + table_prefix: str, + excludable: ExcludableItems, + prefetch_dict: Dict, + orders_by: Dict, ) -> None: """ Instantiates children models extracted from given relation. @@ -610,6 +597,8 @@ class PrefetchQuery: already_extracted dictionary. Later those instances will be fetched by ids and set on the parent model after sorting if needed. + :param excludable: structure of fields to include and exclude + :type excludable: ExcludableItems :param rows: raw sql response from the prefetch query :type rows: List[sqlalchemy.engine.result.RowProxy] :param target_field: field with relation definition from parent model @@ -618,10 +607,6 @@ class PrefetchQuery: :type parent_model: Type[Model] :param table_prefix: prefix of the target table from current relation :type table_prefix: str - :param fields: fields to include - :type fields: Union[Set[Any], Dict[Any, Any], None] - :param exclude_fields: fields to exclude - :type exclude_fields: Union[Set[Any], Dict[Any, Any], None] :param prefetch_dict: dictionaries of related models to prefetch :type prefetch_dict: Dict :param orders_by: dictionary of order by clauses by model @@ -629,16 +614,16 @@ class PrefetchQuery: """ target_model = target_field.to for row in rows: + # TODO Fix fields field_name = parent_model.get_related_field_name(target_field=target_field) item = target_model.extract_prefixed_table_columns( item={}, row=row, table_prefix=table_prefix, - fields=fields, - exclude_fields=exclude_fields, + excludable=excludable, ) item["__excluded__"] = target_model.get_names_to_exclude( - fields=fields, exclude_fields=exclude_fields + excludable=excludable, alias=table_prefix ) instance = target_model(**item) instance = self._populate_nested_related( diff --git a/ormar/queryset/query.py b/ormar/queryset/query.py index d6b10d2..2e88212 100644 --- a/ormar/queryset/query.py +++ b/ormar/queryset/query.py @@ -6,6 +6,7 @@ import sqlalchemy from sqlalchemy import text import ormar # noqa I100 +from ormar.models.excludable import ExcludableItems from ormar.models.helpers.models import group_related_list from ormar.queryset import FilterQuery, LimitQuery, OffsetQuery, OrderQuery from ormar.queryset.actions.filter_action import FilterAction @@ -18,25 +19,23 @@ if TYPE_CHECKING: # pragma no cover class Query: def __init__( # noqa CFQ002 - self, - model_cls: Type["Model"], - filter_clauses: List[FilterAction], - exclude_clauses: List[FilterAction], - select_related: List, - limit_count: Optional[int], - offset: Optional[int], - fields: Optional[Union[Dict, Set]], - exclude_fields: Optional[Union[Dict, Set]], - order_bys: Optional[List["OrderAction"]], - limit_raw_sql: bool, + self, + model_cls: Type["Model"], + filter_clauses: List[FilterAction], + exclude_clauses: List[FilterAction], + select_related: List, + limit_count: Optional[int], + offset: Optional[int], + excludable: ExcludableItems, + order_bys: Optional[List["OrderAction"]], + limit_raw_sql: bool, ) -> None: self.query_offset = offset self.limit_count = limit_count self._select_related = select_related[:] self.filter_clauses = filter_clauses[:] self.exclude_clauses = exclude_clauses[:] - self.fields = copy.deepcopy(fields) if fields else {} - self.exclude_fields = copy.deepcopy(exclude_fields) if exclude_fields else {} + self.excludable = excludable self.model_cls = model_cls self.table = self.model_cls.Meta.table @@ -105,8 +104,7 @@ class Query: """ self_related_fields = self.model_cls.own_table_columns( model=self.model_cls, - fields=self.fields, - exclude_fields=self.exclude_fields, + excludable=self.excludable, use_alias=True, ) self.columns = self.model_cls.Meta.alias_manager.prefixed_columns( @@ -121,8 +119,6 @@ class Query: related_models = group_related_list(self._select_related) for related in related_models: - fields = self.model_cls.get_included(self.fields, related) - exclude_fields = self.model_cls.get_excluded(self.exclude_fields, related) remainder = None if isinstance(related_models, dict) and related_models[related]: remainder = related_models[related] @@ -130,8 +126,7 @@ class Query: used_aliases=self.used_aliases, select_from=self.select_from, columns=self.columns, - fields=fields, - exclude_fields=exclude_fields, + excludable=self.excludable, order_columns=self.order_columns, sorted_orders=self.sorted_orders, main_model=self.model_cls, @@ -196,7 +191,7 @@ class Query: return expr def _apply_expression_modifiers( - self, expr: sqlalchemy.sql.select + self, expr: sqlalchemy.sql.select ) -> sqlalchemy.sql.select: """ Receives the select query (might be join) and applies: @@ -231,5 +226,3 @@ class Query: self.select_from = [] self.columns = [] self.used_aliases = [] - self.fields = {} - self.exclude_fields = {} diff --git a/ormar/queryset/queryset.py b/ormar/queryset/queryset.py index 7c664ac..1c93590 100644 --- a/ormar/queryset/queryset.py +++ b/ormar/queryset/queryset.py @@ -20,6 +20,7 @@ from sqlalchemy import bindparam import ormar # noqa I100 from ormar import MultipleMatches, NoMatch from ormar.exceptions import ModelError, ModelPersistenceError, QueryDefinitionError +from ormar.models.excludable import ExcludableItems from ormar.queryset import FilterQuery from ormar.queryset.actions.order_action import OrderAction from ormar.queryset.clause import QueryClause @@ -41,18 +42,17 @@ class QuerySet(Generic[T]): """ def __init__( # noqa CFQ002 - self, - model_cls: Optional[Type[T]] = None, - filter_clauses: List = None, - exclude_clauses: List = None, - select_related: List = None, - limit_count: int = None, - offset: int = None, - columns: Dict = None, - exclude_columns: Dict = None, - order_bys: List = None, - prefetch_related: List = None, - limit_raw_sql: bool = False, + self, + model_cls: Optional[Type[T]] = None, + filter_clauses: List = None, + exclude_clauses: List = None, + select_related: List = None, + limit_count: int = None, + offset: int = None, + excludable: ExcludableItems = None, + order_bys: List = None, + prefetch_related: List = None, + limit_raw_sql: bool = False, ) -> None: self.model_cls = model_cls self.filter_clauses = [] if filter_clauses is None else filter_clauses @@ -61,15 +61,14 @@ class QuerySet(Generic[T]): self._prefetch_related = [] if prefetch_related is None else prefetch_related self.limit_count = limit_count self.query_offset = offset - self._columns = columns or {} - self._exclude_columns = exclude_columns or {} + self._excludable = excludable or ExcludableItems() self.order_bys = order_bys or [] self.limit_sql_raw = limit_raw_sql def __get__( - self, - instance: Optional[Union["QuerySet", "QuerysetProxy"]], - owner: Union[Type[T], Type["QuerysetProxy"]], + self, + instance: Optional[Union["QuerySet", "QuerysetProxy"]], + owner: Union[Type[T], Type["QuerysetProxy"]], ) -> "QuerySet": if issubclass(owner, ormar.Model): if owner.Meta.requires_ref_update: @@ -107,7 +106,7 @@ class QuerySet(Generic[T]): return self.model_cls async def _prefetch_related_models( - self, models: Sequence[Optional["T"]], rows: List + self, models: Sequence[Optional["T"]], rows: List ) -> Sequence[Optional["T"]]: """ Performs prefetch query for selected models names. @@ -121,8 +120,7 @@ class QuerySet(Generic[T]): """ query = PrefetchQuery( model_cls=self.model, - fields=self._columns, - exclude_fields=self._exclude_columns, + excludable=self._excludable, prefetch_related=self._prefetch_related, select_related=self._select_related, orders_by=self.order_bys, @@ -142,8 +140,7 @@ class QuerySet(Generic[T]): self.model.from_row( row=row, select_related=self._select_related, - fields=self._columns, - exclude_fields=self._exclude_columns, + excludable=self._excludable, source_model=self.model, ) for row in rows @@ -186,7 +183,7 @@ class QuerySet(Generic[T]): return self.model_meta.table def build_select_expression( - self, limit: int = None, offset: int = None, order_bys: List = None, + self, limit: int = None, offset: int = None, order_bys: List = None, ) -> sqlalchemy.sql.select: """ Constructs the actual database query used in the QuerySet. @@ -208,8 +205,7 @@ class QuerySet(Generic[T]): exclude_clauses=self.exclude_clauses, offset=offset or self.query_offset, limit_count=limit or self.limit_count, - fields=self._columns, - exclude_fields=self._exclude_columns, + excludable=self._excludable, order_bys=order_bys or self.order_bys, limit_raw_sql=self.limit_sql_raw, ) @@ -265,8 +261,7 @@ class QuerySet(Generic[T]): select_related=select_related, limit_count=self.limit_count, offset=self.query_offset, - columns=self._columns, - exclude_columns=self._exclude_columns, + excludable=self._excludable, order_bys=self.order_bys, prefetch_related=self._prefetch_related, limit_raw_sql=self.limit_sql_raw, @@ -321,8 +316,7 @@ class QuerySet(Generic[T]): select_related=related, limit_count=self.limit_count, offset=self.query_offset, - columns=self._columns, - exclude_columns=self._exclude_columns, + excludable=self._excludable, order_bys=self.order_bys, prefetch_related=self._prefetch_related, limit_raw_sql=self.limit_sql_raw, @@ -357,14 +351,14 @@ class QuerySet(Generic[T]): select_related=self._select_related, limit_count=self.limit_count, offset=self.query_offset, - columns=self._columns, - exclude_columns=self._exclude_columns, + excludable=self._excludable, order_bys=self.order_bys, prefetch_related=related, limit_raw_sql=self.limit_sql_raw, ) - def fields(self, columns: Union[List, str, Set, Dict]) -> "QuerySet": + def fields(self, columns: Union[List, str, Set, Dict], + _is_exclude: bool = False) -> "QuerySet": """ With `fields()` you can select subset of model columns to limit the data load. @@ -407,15 +401,10 @@ class QuerySet(Generic[T]): :return: QuerySet :rtype: QuerySet """ - if isinstance(columns, str): - columns = [columns] - - # TODO: Flatten all excludes into one dict-like structure with alias + model key - current_included = self._columns - if not isinstance(columns, dict): - current_included = update_dict_from_list(current_included, columns) - else: - current_included = update(current_included, columns) + excludable = ExcludableItems.from_excludable(self._excludable) + excludable.build(items=columns, + model_cls=self.model_cls, + is_exclude=_is_exclude) return self.__class__( model_cls=self.model, @@ -424,8 +413,7 @@ class QuerySet(Generic[T]): select_related=self._select_related, limit_count=self.limit_count, offset=self.query_offset, - columns=current_included, - exclude_columns=self._exclude_columns, + excludable=excludable, order_bys=self.order_bys, prefetch_related=self._prefetch_related, limit_raw_sql=self.limit_sql_raw, @@ -458,28 +446,7 @@ class QuerySet(Generic[T]): :return: QuerySet :rtype: QuerySet """ - if isinstance(columns, str): - columns = [columns] - - current_excluded = self._exclude_columns - if not isinstance(columns, dict): - current_excluded = update_dict_from_list(current_excluded, columns) - else: - current_excluded = update(current_excluded, columns) - - return self.__class__( - model_cls=self.model, - filter_clauses=self.filter_clauses, - exclude_clauses=self.exclude_clauses, - select_related=self._select_related, - limit_count=self.limit_count, - offset=self.query_offset, - columns=self._columns, - exclude_columns=current_excluded, - order_bys=self.order_bys, - prefetch_related=self._prefetch_related, - limit_raw_sql=self.limit_sql_raw, - ) + return self.fields(columns=columns, _is_exclude=True) def order_by(self, columns: Union[List, str]) -> "QuerySet": """ @@ -529,8 +496,7 @@ class QuerySet(Generic[T]): select_related=self._select_related, limit_count=self.limit_count, offset=self.query_offset, - columns=self._columns, - exclude_columns=self._exclude_columns, + excludable=self._excludable, order_bys=order_bys, prefetch_related=self._prefetch_related, limit_raw_sql=self.limit_sql_raw, @@ -642,8 +608,7 @@ class QuerySet(Generic[T]): select_related=self._select_related, limit_count=limit_count, offset=query_offset, - columns=self._columns, - exclude_columns=self._exclude_columns, + excludable=self._excludable, order_bys=self.order_bys, prefetch_related=self._prefetch_related, limit_raw_sql=self.limit_sql_raw, @@ -671,8 +636,7 @@ class QuerySet(Generic[T]): select_related=self._select_related, limit_count=limit_count, offset=self.query_offset, - columns=self._columns, - exclude_columns=self._exclude_columns, + excludable=self._excludable, order_bys=self.order_bys, prefetch_related=self._prefetch_related, limit_raw_sql=limit_raw_sql, @@ -700,8 +664,7 @@ class QuerySet(Generic[T]): select_related=self._select_related, limit_count=self.limit_count, offset=offset, - columns=self._columns, - exclude_columns=self._exclude_columns, + excludable=self._excludable, order_bys=self.order_bys, prefetch_related=self._prefetch_related, limit_raw_sql=limit_raw_sql, @@ -724,12 +687,12 @@ class QuerySet(Generic[T]): expr = self.build_select_expression( limit=1, order_bys=[ - OrderAction( - order_str=f"{self.model.Meta.pkname}", - model_cls=self.model_cls, # type: ignore - ) - ] - + self.order_bys, + OrderAction( + order_str=f"{self.model.Meta.pkname}", + model_cls=self.model_cls, # type: ignore + ) + ] + + self.order_bys, ) rows = await self.database.fetch_all(expr) processed_rows = self._process_query_result_rows(rows) @@ -760,12 +723,12 @@ class QuerySet(Generic[T]): expr = self.build_select_expression( limit=1, order_bys=[ - OrderAction( - order_str=f"-{self.model.Meta.pkname}", - model_cls=self.model_cls, # type: ignore - ) - ] - + self.order_bys, + OrderAction( + order_str=f"-{self.model.Meta.pkname}", + model_cls=self.model_cls, # type: ignore + ) + ] + + self.order_bys, ) else: expr = self.build_select_expression() @@ -868,9 +831,9 @@ class QuerySet(Generic[T]): # refresh server side defaults if any( - field.server_default is not None - for name, field in self.model.Meta.model_fields.items() - if name not in kwargs + field.server_default is not None + for name, field in self.model.Meta.model_fields.items() + if name not in kwargs ): instance = await instance.load() instance.set_save_status(True) @@ -905,7 +868,7 @@ class QuerySet(Generic[T]): objt.set_save_status(True) async def bulk_update( # noqa: CCR001 - self, objects: List[T], columns: List[str] = None + self, objects: List[T], columns: List[str] = None ) -> None: """ Performs bulk update in one database session to speed up the process.