change limit/offset with select related to be applied on a subquery and limit only main model query
This commit is contained in:
@ -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:
|
||||
|
||||
@ -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():
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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}))
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
Reference in New Issue
Block a user