change limit/offset with select related to be applied on a subquery and limit only main model query

This commit is contained in:
collerek
2020-12-21 18:42:17 +01:00
parent 514e8c4ad5
commit c8aad2385c
18 changed files with 329 additions and 69 deletions

View File

@ -48,8 +48,8 @@ def register_many_to_many_relation_on_build(
:param field: relation field
:type field: ManyToManyField class
"""
alias_manager.add_relation_type(field.through, new_model.get_name(), is_multi=True)
alias_manager.add_relation_type(field.through, field.to.get_name(), is_multi=True)
alias_manager.add_relation_type(field.through, new_model.get_name())
alias_manager.add_relation_type(field.through, field.to.get_name())
def expand_reverse_relationships(model: Type["Model"]) -> None:

View File

@ -17,13 +17,12 @@ from typing import (
Union,
)
from ormar.exceptions import ModelPersistenceError, RelationshipInstanceError
from ormar.queryset.utils import translate_list_to_dict, update
import ormar # noqa: I100
from ormar.exceptions import ModelPersistenceError
from ormar.fields import BaseField, ManyToManyField
from ormar.fields.foreign_key import ForeignKeyField
from ormar.models.metaclass import ModelMeta
from ormar.queryset.utils import translate_list_to_dict, update
if TYPE_CHECKING: # pragma no cover
from ormar import Model
@ -76,9 +75,10 @@ class ModelTableProxy:
)
field = target_model.Meta.model_fields[field_name]
if issubclass(field, ormar.fields.ManyToManyField):
sub_field = target_model.resolve_relation_field(
field.through, parent_model
field_name = parent_model.resolve_relation_name(
field.through, field.to, explicit_multi=True
)
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)
@ -86,17 +86,14 @@ class ModelTableProxy:
@staticmethod
def get_column_name_for_id_extraction(
parent_model: Type["Model"],
target_model: Type["Model"],
reverse: bool,
use_raw: bool,
parent_model: Type["Model"], reverse: bool, related: str, use_raw: bool,
) -> str:
if reverse:
column_name = parent_model.Meta.pkname
return (
parent_model.get_column_alias(column_name) if use_raw else column_name
)
column = target_model.resolve_relation_field(parent_model, target_model)
column = parent_model.Meta.model_fields[related]
return column.get_alias() if use_raw else column.name
@classmethod
@ -322,19 +319,6 @@ class ModelTableProxy:
f"No relation between {item.get_name()} and {related.get_name()}"
) # pragma nocover
@staticmethod
def resolve_relation_field(
item: Union["Model", Type["Model"]], related: Union["Model", Type["Model"]]
) -> Type[BaseField]:
name = ModelTableProxy.resolve_relation_name(item, related)
to_field = item.Meta.model_fields.get(name)
if not to_field: # pragma no cover
raise RelationshipInstanceError(
f"Model {item.__class__} does not have "
f"reference to model {related.__class__}"
)
return to_field
@classmethod
def translate_columns_to_aliases(cls, new_kwargs: Dict) -> Dict:
for field_name, field in cls.Meta.model_fields.items():

View File

@ -224,8 +224,8 @@ class NewBaseModel(
def db_backend_name(cls) -> str:
return cls.Meta.database._backend._dialect.name
def remove(self, name: "T") -> None:
self._orm.remove_parent(self, name)
def remove(self, parent: "T", name: str) -> None:
self._orm.remove_parent(self, parent, name)
def set_save_status(self, status: bool) -> None:
object.__setattr__(self, "_orm_saved", status)

View File

@ -138,7 +138,7 @@ class QueryClause:
through_field = model_cls.Meta.model_fields[part]
previous_model = through_field.through
part2 = model_cls.resolve_relation_name(
through_field.through, through_field.to, explicit_multi=True
previous_model, through_field.to, explicit_multi=True
)
manager = model_cls.Meta.alias_manager
table_prefix = manager.resolve_relation_join(previous_model, part2)

View File

@ -123,15 +123,15 @@ class PrefetchQuery:
return list_of_ids
def _extract_required_ids(
self, parent_model: Type["Model"], target_model: Type["Model"], reverse: bool,
self, parent_model: Type["Model"], reverse: bool, related: str,
) -> Set:
use_raw = parent_model.get_name() not in self.models
column_name = parent_model.get_column_name_for_id_extraction(
parent_model=parent_model,
target_model=target_model,
reverse=reverse,
related=related,
use_raw=use_raw,
)
@ -152,7 +152,7 @@ class PrefetchQuery:
related: str,
) -> List:
ids = self._extract_required_ids(
parent_model=parent_model, target_model=target_model, reverse=reverse,
parent_model=parent_model, reverse=reverse, related=related
)
if ids:
(
@ -343,6 +343,7 @@ class PrefetchQuery:
fields=fields,
exclude_fields=exclude_fields,
order_bys=None,
limit_raw_sql=False,
)
expr = qry.build_select_expression()
# print(expr.compile(compile_kwargs={"literal_binds": True}))

View File

@ -25,6 +25,7 @@ class Query:
fields: Optional[Union[Dict, Set]],
exclude_fields: Optional[Union[Dict, Set]],
order_bys: Optional[List],
limit_raw_sql: bool,
) -> None:
self.query_offset = offset
self.limit_count = limit_count
@ -45,6 +46,8 @@ class Query:
self.sorted_orders: OrderedDict = OrderedDict()
self._init_sorted_orders()
self.limit_raw_sql = limit_raw_sql
def _init_sorted_orders(self) -> None:
if self.order_columns:
for clause in self.order_columns:
@ -62,16 +65,31 @@ class Query:
if self.order_columns:
for clause in self.order_columns:
if "__" not in clause:
clause = (
text_clause = (
text(f"{self.alias(clause[1:])} desc")
if clause.startswith("-")
else text(self.alias(clause))
)
self.sorted_orders[clause] = clause
self.sorted_orders[clause] = text_clause
else:
order = text(self.prefixed_pk_name)
self.sorted_orders[self.prefixed_pk_name] = order
def _pagination_query_required(self) -> bool:
"""
Checks if limit or offset are set, the flag limit_sql_raw is not set
and query has select_related applied. Otherwise we can limit/offset normally
at the end of whole query.
:return: result of the check
:rtype: bool
"""
return bool(
(self.limit_count or self.query_offset)
and not self.limit_raw_sql
and self._select_related
)
def build_select_expression(self) -> Tuple[sqlalchemy.sql.select, List[str]]:
self_related_fields = self.model_cls.own_table_columns(
model=self.model_cls,
@ -83,7 +101,10 @@ class Query:
"", self.table, self_related_fields
)
self.apply_order_bys_for_primary_model()
self.select_from = self.table
if self._pagination_query_required():
self.select_from = self._build_pagination_subquery()
else:
self.select_from = self.table
self._select_related.sort(key=lambda item: (item, -len(item)))
@ -120,6 +141,46 @@ class Query:
return expr
def _build_pagination_subquery(self) -> sqlalchemy.sql.select:
"""
In order to apply limit and offset on main table in join only
(otherwise you can get only partially constructed main model
if number of children exceeds the applied limit and select_related is used)
Used also to change first and get() without argument behaviour.
Needed only if limit or offset are set, the flag limit_sql_raw is not set
and query has select_related applied. Otherwise we can limit/offset normally
at the end of whole query.
:return: constructed subquery on main table with limit, offset and order applied
:rtype: sqlalchemy.sql.select
"""
expr = sqlalchemy.sql.select(self.model_cls.Meta.table.columns)
expr = LimitQuery(limit_count=self.limit_count).apply(expr)
expr = OffsetQuery(query_offset=self.query_offset).apply(expr)
filters_to_use = [
filter_clause
for filter_clause in self.filter_clauses
if filter_clause.text.startswith(f"{self.table.name}.")
]
excludes_to_use = [
filter_clause
for filter_clause in self.exclude_clauses
if filter_clause.text.startswith(f"{self.table.name}.")
]
sorts_to_use = {
k: v
for k, v in self.sorted_orders.items()
if k.startswith(f"{self.table.name}.")
}
expr = FilterQuery(filter_clauses=filters_to_use).apply(expr)
expr = FilterQuery(filter_clauses=excludes_to_use, exclude=True).apply(expr)
expr = OrderQuery(sorted_orders=sorts_to_use).apply(expr)
expr = expr.alias(f"{self.table}")
self.filter_clauses = list(set(self.filter_clauses) - set(filters_to_use))
self.exclude_clauses = list(set(self.exclude_clauses) - set(excludes_to_use))
return expr
def _apply_expression_modifiers(
self, expr: sqlalchemy.sql.select
) -> sqlalchemy.sql.select:
@ -127,8 +188,9 @@ class Query:
expr = FilterQuery(filter_clauses=self.exclude_clauses, exclude=True).apply(
expr
)
expr = LimitQuery(limit_count=self.limit_count).apply(expr)
expr = OffsetQuery(query_offset=self.query_offset).apply(expr)
if not self._pagination_query_required():
expr = LimitQuery(limit_count=self.limit_count).apply(expr)
expr = OffsetQuery(query_offset=self.query_offset).apply(expr)
expr = OrderQuery(sorted_orders=self.sorted_orders).apply(expr)
return expr

View File

@ -32,6 +32,7 @@ class QuerySet:
exclude_columns: Dict = 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
@ -43,6 +44,7 @@ class QuerySet:
self._columns = columns or {}
self._exclude_columns = exclude_columns or {}
self.order_bys = order_bys or []
self.limit_sql_raw = limit_raw_sql
def __get__(
self,
@ -123,17 +125,20 @@ class QuerySet:
def table(self) -> sqlalchemy.Table:
return self.model_meta.table
def build_select_expression(self) -> sqlalchemy.sql.select:
def build_select_expression(
self, limit: int = None, offset: int = None, order_bys: List = None,
) -> sqlalchemy.sql.select:
qry = Query(
model_cls=self.model,
select_related=self._select_related,
filter_clauses=self.filter_clauses,
exclude_clauses=self.exclude_clauses,
offset=self.query_offset,
limit_count=self.limit_count,
offset=offset or self.query_offset,
limit_count=limit or self.limit_count,
fields=self._columns,
exclude_fields=self._exclude_columns,
order_bys=self.order_bys,
order_bys=order_bys or self.order_bys,
limit_raw_sql=self.limit_sql_raw,
)
exp = qry.build_select_expression()
# print(exp.compile(compile_kwargs={"literal_binds": True}))
@ -164,6 +169,7 @@ class QuerySet:
exclude_columns=self._exclude_columns,
order_bys=self.order_bys,
prefetch_related=self._prefetch_related,
limit_raw_sql=self.limit_sql_raw,
)
def exclude(self, **kwargs: Any) -> "QuerySet": # noqa: A003
@ -185,6 +191,7 @@ class QuerySet:
exclude_columns=self._exclude_columns,
order_bys=self.order_bys,
prefetch_related=self._prefetch_related,
limit_raw_sql=self.limit_sql_raw,
)
def prefetch_related(self, related: Union[List, str]) -> "QuerySet":
@ -203,6 +210,7 @@ class QuerySet:
exclude_columns=self._exclude_columns,
order_bys=self.order_bys,
prefetch_related=related,
limit_raw_sql=self.limit_sql_raw,
)
def exclude_fields(self, columns: Union[List, str, Set, Dict]) -> "QuerySet":
@ -226,6 +234,7 @@ class QuerySet:
exclude_columns=current_excluded,
order_bys=self.order_bys,
prefetch_related=self._prefetch_related,
limit_raw_sql=self.limit_sql_raw,
)
def fields(self, columns: Union[List, str, Set, Dict]) -> "QuerySet":
@ -249,6 +258,7 @@ class QuerySet:
exclude_columns=self._exclude_columns,
order_bys=self.order_bys,
prefetch_related=self._prefetch_related,
limit_raw_sql=self.limit_sql_raw,
)
def order_by(self, columns: Union[List, str]) -> "QuerySet":
@ -267,6 +277,7 @@ class QuerySet:
exclude_columns=self._exclude_columns,
order_bys=order_bys,
prefetch_related=self._prefetch_related,
limit_raw_sql=self.limit_sql_raw,
)
async def exists(self) -> bool:
@ -308,7 +319,8 @@ class QuerySet:
)
return await self.database.execute(expr)
def limit(self, limit_count: int) -> "QuerySet":
def limit(self, limit_count: int, limit_raw_sql: bool = None) -> "QuerySet":
limit_raw_sql = self.limit_sql_raw if limit_raw_sql is None else limit_raw_sql
return self.__class__(
model_cls=self.model,
filter_clauses=self.filter_clauses,
@ -320,9 +332,11 @@ class QuerySet:
exclude_columns=self._exclude_columns,
order_bys=self.order_bys,
prefetch_related=self._prefetch_related,
limit_raw_sql=limit_raw_sql,
)
def offset(self, offset: int) -> "QuerySet":
def offset(self, offset: int, limit_raw_sql: bool = None) -> "QuerySet":
limit_raw_sql = self.limit_sql_raw if limit_raw_sql is None else limit_raw_sql
return self.__class__(
model_cls=self.model,
filter_clauses=self.filter_clauses,
@ -334,23 +348,33 @@ class QuerySet:
exclude_columns=self._exclude_columns,
order_bys=self.order_bys,
prefetch_related=self._prefetch_related,
limit_raw_sql=limit_raw_sql,
)
async def first(self, **kwargs: Any) -> "Model":
if kwargs:
return await self.filter(**kwargs).first()
rows = await self.limit(1).all()
self.check_single_result_rows_count(rows)
return rows[0] # type: ignore
expr = self.build_select_expression(
limit=1, order_bys=[f"{self.model.Meta.pkname}"]
)
rows = await self.database.fetch_all(expr)
processed_rows = self._process_query_result_rows(rows)
if self._prefetch_related and processed_rows:
processed_rows = await self._prefetch_related_models(processed_rows, rows)
self.check_single_result_rows_count(processed_rows)
return processed_rows[0] # type: ignore
async def get(self, **kwargs: Any) -> "Model":
if kwargs:
return await self.filter(**kwargs).get()
expr = self.build_select_expression()
if not self.filter_clauses:
expr = expr.limit(2)
expr = self.build_select_expression(
limit=1, order_bys=[f"-{self.model.Meta.pkname}"]
)
else:
expr = self.build_select_expression()
rows = await self.database.fetch_all(expr)
processed_rows = self._process_query_result_rows(rows)

View File

@ -40,7 +40,7 @@ class AliasManager:
return text(f"{name} {alias}_{name}")
def add_relation_type(
self, source_model: Type["Model"], relation_name: str, is_multi: bool = False
self, source_model: Type["Model"], relation_name: str
) -> None:
parent_key = f"{source_model.get_name()}_{relation_name}"
if parent_key not in self._aliases_new:
@ -50,7 +50,7 @@ class AliasManager:
related_name = to_field.related_name
if not related_name:
related_name = child_model.resolve_relation_name(
child_model, source_model, explicit_multi=is_multi
child_model, source_model, explicit_multi=True
)
child_key = f"{child_model.get_name()}_{related_name}"
if child_key not in self._aliases_new:

View File

@ -53,7 +53,7 @@ class QuerysetProxy(ormar.QuerySetProtocol):
def _assign_child_to_parent(self, child: Optional["T"]) -> None:
if child:
owner = self._owner
rel_name = owner.resolve_relation_name(owner, child)
rel_name = self.relation.field_name
setattr(owner, rel_name, child)
def _register_related(self, child: Union["T", Sequence[Optional["T"]]]) -> None:

View File

@ -89,11 +89,10 @@ class RelationsManager:
@staticmethod
def remove_parent(
item: Union["NewBaseModel", Type["NewBaseModel"]], name: "Model"
item: Union["NewBaseModel", Type["NewBaseModel"]], parent: "Model", name: str
) -> None:
related_model = name
rel_name = item.resolve_relation_name(item, related_model)
if rel_name in item._orm:
relation_name = item.resolve_relation_name(related_model, item)
item._orm.remove(rel_name, related_model)
related_model._orm.remove(relation_name, item)
relation_name = (
item.Meta.model_fields[name].related_name or item.get_name() + "s"
)
item._orm.remove(name, parent)
parent._orm.remove(relation_name, item)