diff --git a/README.md b/README.md index 939b7c3a3..ad044eace 100644 --- a/README.md +++ b/README.md @@ -620,6 +620,7 @@ asyncio.run(cleanup_database()) * `min(columns: list[str]) -> Any` * `avg(columns: list[str]) -> Any` * `sum(columns: list[str]) -> Any` +* `annotate(**aggregates) -> QuerySet` / `having(**filters) -> QuerySet` - per-parent `Count`/`Sum`/`Avg`/`Min`/`Max` annotations and filtering on them * `fields(columns: Union[list, str, set, dict]) -> QuerySet` * `exclude_fields(columns: Union[list, str, set, dict]) -> QuerySet` * `order_by(columns:Union[list, str]) -> QuerySet` diff --git a/docs/index.md b/docs/index.md index 97959527c..0d50ad30a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -625,6 +625,7 @@ asyncio.run(cleanup_database()) * `min(columns: list[str]) -> Any` * `avg(columns: list[str]) -> Any` * `sum(columns: list[str]) -> Any` +* `annotate(**aggregates) -> QuerySet` / `having(**filters) -> QuerySet` - per-parent `Count`/`Sum`/`Avg`/`Min`/`Max` annotations and filtering on them * `fields(columns: Union[list, str, set, dict]) -> QuerySet` * `exclude_fields(columns: Union[list, str, set, dict]) -> QuerySet` * `order_by(columns:Union[list, str]) -> QuerySet` diff --git a/docs/queries/aggregate-annotations.md b/docs/queries/aggregate-annotations.md new file mode 100644 index 000000000..0a1278dfe --- /dev/null +++ b/docs/queries/aggregate-annotations.md @@ -0,0 +1,242 @@ +# Aggregate annotations + +`annotate()` adds a named, per-parent aggregate value (count/sum/avg/min/max of a +related model) to a query, computed with a single pre-grouped derived-table +`LEFT JOIN` rather than a correlated subquery per row. It composes with +`filter()`, `having()`, `order_by()`, `select_related()`, `limit()`/`offset()`, +and `values()`/`values_list()`. + +* `annotate(**aggregates: AggregateFunction) -> QuerySet` +* `having(**filters: Any) -> QuerySet` + +Aggregate expression objects (imported from the `ormar` top level): + +* `ormar.Count(field: str, distinct: bool = False)` +* `ormar.Sum(field: str)` +* `ormar.Avg(field: str)` +* `ormar.Min(field: str)` +* `ormar.Max(field: str)` + +!!!note + This page documents phase 1 of aggregation support ("Mode A" - one aggregate + value per parent row). True `GROUP BY` rollups over a non-unique column + (`.group_by("country")`) are not implemented yet - see + [Limitations](#limitations) below. + +## annotate + +`field` is either a relation name (`"tasks"`) for a bare `Count`, or a +`relation__column` path (`"tasks__price"`) for `Sum`/`Avg`/`Min`/`Max` (and +optionally `Count` too). + +Given the following models: + +```python +class User(ormar.Model): + ormar_config = base_ormar_config.copy(tablename="users") + + id: int = ormar.Integer(primary_key=True) + name: str = ormar.String(max_length=100) + + +class Task(ormar.Model): + ormar_config = base_ormar_config.copy(tablename="tasks") + + id: int = ormar.Integer(primary_key=True) + user: Optional[User] = ormar.ForeignKey(User, related_name="tasks") + title: str = ormar.String(max_length=100) + price: Optional[int] = ormar.Integer(nullable=True) +``` + +```python +rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("name") + .values(["name", "task_count"]) +) +assert rows == [ + {"name": "Alice", "task_count": 2}, + {"name": "Bob", "task_count": 1}, + {"name": "Carol", "task_count": 0}, # no tasks +] +``` + +Several aggregates can be requested at once, including several columns of the +same relation: + +```python +rows = ( + await User.objects.annotate( + total=ormar.Sum("tasks__price"), + avg_price=ormar.Avg("tasks__price"), + cheapest=ormar.Min("tasks__price"), + dearest=ormar.Max("tasks__price"), + ) + .filter(name="Alice") + .values(["name", "total", "cheapest", "dearest"]) +) +assert rows == [{"name": "Alice", "total": 30, "cheapest": 10, "dearest": 20}] +``` + +`Count(field, distinct=True)` counts distinct related rows/values, matching the +existing `distinct` flag on `QuerySet.count()`: + +```python +await User.objects.annotate( + task_count=ormar.Count("tasks__id", distinct=True) +).values(["name", "task_count"]) +``` + +Passing `distinct=True` to a bare `Count("relation")` (no column) is a no-op: +distinct only changes anything when it is applied to an explicit column, so use +`Count("relation__column", distinct=True)` as shown above if you need it. + +### Zero vs. `NULL` for parents with no children + +`Count` over a parent with no related rows returns `0` (the derived aggregate is +wrapped in `COALESCE(..., 0)`). `Sum`/`Avg`/`Min`/`Max` are **not** coalesced, so +they return `None` for a childless parent - there is no sensible zero for "the +average of no rows": + +```python +carol = ( + await User.objects.annotate(total=ormar.Sum("tasks__price")) + .filter(name="Carol") # Carol has no tasks + .values(["name", "total"]) +) +assert carol == [{"name": "Carol", "total": None}] +``` + +## Ordering by an annotation + +`order_by()` accepts an annotated name exactly like a regular column, ascending +or descending (this closes +[issue #566](https://github.com/ormar-orm/ormar/issues/566), sorting parents by +the number of related rows): + +```python +names = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("-task_count") + .values_list("name", flatten=True) +) +assert names == ["Alice", "Bob", "Carol"] +``` + +## having + +`having()` filters rows by an annotated aggregate value - the annotation +equivalent of `filter()`. Supported suffixes are `exact` (default), `gt`, +`gte`, `lt`, `lte` and `ne`: + +```python +rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .having(task_count__gt=1) + .values_list("name", flatten=True) +) +assert rows == ["Alice"] +``` + +Any other suffix (e.g. `contains`) raises `QueryDefinitionError`. + +## values() and values_list() output + +Annotated names flow through `values()`/`values_list()` as extra keys/columns, +whether or not you pass an explicit field subset. If you don't pass `fields`, +every regular column is returned alongside every annotation: + +```python +rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("name") + .values() +) +# every row carries the normal columns (id, name, ...) plus task_count +assert all("task_count" in row for row in rows) +``` + +Passing an explicit field list works the same way as for regular columns - +list the annotation name alongside the columns you want: + +```python +await User.objects.annotate(task_count=ormar.Count("tasks")).values( + ["name", "task_count"] +) +``` + +!!!note + Annotated values are only available through `values()`/`values_list()` in + this phase. They are not (yet) attached as attributes on hydrated model + instances returned by `all()`/`get()`. + +## Many-to-many relations + +A bare `Count` over a many-to-many relation is supported - it groups over the +through table, so the far side of the relation is never joined: + +```python +rows = ( + await Post.objects.annotate(tag_count=ormar.Count("tags")) + .order_by("title") + .values(["title", "tag_count"]) +) +``` + +A **qualified** aggregate over a many-to-many relation (`Sum("tags__price")`, +or even `Count("tags__id")`) is not supported and raises `QueryDefinitionError`. +Only `Count("relation")` without a column works over many-to-many. + +## Composing with select_related and pagination + +Because the aggregate is joined 1:1 onto the parent primary key, it does not +change the number of parent rows, so it composes freely with `select_related()` +and with `limit()`/`offset()`: + +```python +users = ( + await User.objects.select_related("tasks") + .annotate(task_count=ormar.Count("tasks")) + .order_by("-task_count") + .limit(2) + .all() +) +``` + +## Limitations + +- **No `group_by()` / rollups yet.** This phase only supports one aggregate + value per parent row (joined by primary key). Grouping by an arbitrary, + non-unique column (e.g. "count of tasks per country") is the more general + request tracked in + [issue #266](https://github.com/ormar-orm/ormar/issues/266) and is planned as + a future addition, not covered by `annotate()`/`having()` today. +- **Many-to-many is Count-only.** A qualified aggregate over a many-to-many + relation (`Sum("tags__price")`, `Count("tags__id")`) raises + `QueryDefinitionError`; only bare `Count("relation")` is supported over + many-to-many, as shown above. +- **An outer to-many relation filter can duplicate parent rows.** This is + pre-existing `ormar` behaviour unrelated to annotations (see the `distinct` + flag on `QuerySet.count()`): filtering on a to-many relation's column + (`.filter(tasks__price__gt=15)`) joins that relation again to evaluate the + filter, and plain SQL join semantics duplicate the parent row once per + matching child in the flat `values()`/`values_list()` result. The annotation + itself is unaffected by that outer filter - it is computed by its own + derived-table subquery grouped over the raw child table, so it reports the + same, correct value (based on **all** matching children, not just the ones + matched by the outer filter) on every duplicated row: + +```python +rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .filter(tasks__price__gt=15) # matches 2 of Alice's 3 tasks + .order_by("name") + .values(["name", "task_count"]) +) +# Alice's row is duplicated by the outer filter's join (pre-existing +# behaviour), but task_count is 3 (all of Alice's tasks) on both rows. +assert rows == [ + {"name": "Alice", "task_count": 3}, + {"name": "Alice", "task_count": 3}, +] +``` diff --git a/mkdocs.yml b/mkdocs.yml index 558c6815c..91eb1839d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -32,6 +32,7 @@ nav: - queries/select-columns.md - queries/pagination-and-rows-number.md - queries/aggregations.md + - queries/aggregate-annotations.md - Return raw data: queries/raw-data.md - Signals: signals.md - Transactions: transactions.md diff --git a/ormar/__init__.py b/ormar/__init__.py index 319178853..7ad89758e 100644 --- a/ormar/__init__.py +++ b/ormar/__init__.py @@ -76,6 +76,7 @@ from ormar.databases.connection import DatabaseConnection from ormar.models import ExcludableItems, Extra, Model, OrmarConfig from ormar.queryset import NullsOrdering, OrderAction, QuerySet, and_, or_ +from ormar.queryset.aggregations import Avg, Count, Max, Min, Sum from ormar.relations import RelationType from ormar.signals import Signal @@ -109,6 +110,11 @@ def __repr__(self) -> str: "NoMatch", "ForeignKey", "QuerySet", + "Count", + "Sum", + "Avg", + "Min", + "Max", "RelationType", "Undefined", "UUID", diff --git a/ormar/queryset/actions/aggregation_action.py b/ormar/queryset/actions/aggregation_action.py new file mode 100644 index 000000000..109dd8f82 --- /dev/null +++ b/ormar/queryset/actions/aggregation_action.py @@ -0,0 +1,154 @@ +"""Builds the derived-table join that backs a single ``annotate`` aggregate.""" + +from typing import TYPE_CHECKING, Optional, cast + +import sqlalchemy + +from ormar.exceptions import QueryDefinitionError +from ormar.queryset.aggregations import AggregateFunction +from ormar.queryset.utils import get_relationship_alias_model_and_str + +if TYPE_CHECKING: # pragma: no cover + from ormar import ManyToManyField, Model + + +class AggregationAction: + """ + Compiles one ``annotate(name=Func("relation[__column]"))`` entry into a + pre-grouped derived table joined 1:1 onto the parent by primary key. + + :param name: label of the annotation in the result set + :type name: str + :param aggregate: the requested aggregate function value object + :type aggregate: AggregateFunction + :param model_cls: the queried (parent) model + :type model_cls: type["Model"] + """ + + def __init__( + self, name: str, aggregate: AggregateFunction, model_cls: type["Model"] + ) -> None: + self.name = name + self.aggregate = aggregate + self.source_model = model_cls + parts = aggregate.field.split("__") + self.relation_name = parts[0] + self.column_name: Optional[str] = parts[1] if len(parts) > 1 else None + _, self.target_model, _, _ = get_relationship_alias_model_and_str( + model_cls, [self.relation_name] + ) + self.result_column: sqlalchemy.sql.ColumnElement = None # type: ignore + + def _child_group_key(self) -> sqlalchemy.Column: + """ + Returns the child-table FK column pointing back to the parent. + + Mirrors ``SqlJoin._get_to_and_from_keys`` (``ormar/queryset/join.py``) + for the reverse-FK (``virtual``) branch: the relation field stored on + the parent model exposes ``get_related_name()``, which resolves to + the name of the FK field declared on the child model. That name is + then translated to its database column alias on the child table. + + :return: child table column used as the ``GROUP BY`` key + :rtype: sqlalchemy.Column + """ + relation_field = self.source_model.ormar_config.model_fields[self.relation_name] + related_name = relation_field.get_related_name() + fk_alias = self.target_model.get_column_alias(related_name) + return self.target_model.ormar_config.table.columns[fk_alias] + + def _is_m2m(self) -> bool: + """ + Returns whether the annotated relation is a Many-to-Many relation. + + :return: ``True`` if the relation field is a ``ManyToManyField`` + :rtype: bool + """ + return bool( + self.source_model.ormar_config.model_fields[self.relation_name].is_multi + ) + + def _m2m_group_key(self) -> sqlalchemy.Column: + """ + Returns the through-table FK column pointing back to the parent. + + Mirrors ``SqlJoin._get_to_and_from_keys`` (``ormar/queryset/join.py``) + for the ``is_multi`` branch: the relation field exposes its ``through`` + model and ``default_source_field_name()``, which resolves to the name + of the FK field on the through model pointing back to the owner + (parent) model. That name is then translated to its database column + alias on the through table, so counting can be grouped per parent + without joining the far side of the relation at all. + + :return: through table column used as the ``GROUP BY`` key + :rtype: sqlalchemy.Column + """ + field = cast( + "ManyToManyField", + self.source_model.ormar_config.model_fields[self.relation_name], + ) + through = field.through + source_name = field.default_source_field_name() + fk_alias = through.get_column_alias(source_name) + return through.ormar_config.table.columns[fk_alias] + + def _aggregate_target(self) -> sqlalchemy.sql.ColumnElement: + """ + Returns the column the aggregate function is applied to. + + :return: ``*`` literal for count-all, otherwise the resolved column + :rtype: sqlalchemy.sql.ColumnElement + """ + if self.column_name is None: + return sqlalchemy.literal_column("*") + col_alias = self.target_model.get_column_alias(self.column_name) + return self.target_model.ormar_config.table.columns[col_alias] + + def apply_join( + self, + select_from: sqlalchemy.sql.expression.FromClause, + parent_table: sqlalchemy.Table, + ) -> sqlalchemy.sql.expression.FromClause: + """ + Builds the grouped derived table, LEFT JOINs it to ``select_from`` on the + parent primary key, stores the labelled result column and returns the new + from-clause. + + :param select_from: current from-clause the join is appended to + :type select_from: sqlalchemy.sql.expression.FromClause + :param parent_table: table (or alias) of the queried (parent) model + :type parent_table: sqlalchemy.Table + :return: from-clause extended with the derived-table LEFT JOIN + :rtype: sqlalchemy.sql.expression.FromClause + :raises QueryDefinitionError: when a specific column is aggregated + across a many-to-many relation, which is not supported + """ + is_m2m = self._is_m2m() + if is_m2m and self.column_name is not None: + raise QueryDefinitionError( + "Aggregating a specific column across a many-to-many relation " + f"is not supported (relation '{self.relation_name}', column " + f"'{self.column_name}'). Only Count('{self.relation_name}') " + "over a many-to-many relation is supported." + ) + group_key = self._m2m_group_key() if is_m2m else self._child_group_key() + target = self._aggregate_target() + func = getattr(sqlalchemy.func, self.aggregate.function_name) + use_distinct = self.aggregate.distinct and self.column_name is not None + expr = func(target.distinct()) if use_distinct else func(target) + derived = ( + sqlalchemy.select(group_key.label("ormar_agg_key"), expr.label(self.name)) + .group_by(group_key) + .alias(f"{self.name}_agg") + ) + pk_alias = self.source_model.get_column_alias( + self.source_model.ormar_config.pkname + ) + parent_pk = parent_table.columns[pk_alias] + value: sqlalchemy.sql.ColumnElement = derived.c[self.name] + if self.aggregate.function_name == "count": + value = sqlalchemy.func.coalesce(value, 0) + self.result_column = value.label(self.name) + return sqlalchemy.sql.outerjoin( + select_from, derived, derived.c.ormar_agg_key == parent_pk + ) diff --git a/ormar/queryset/aggregations.py b/ormar/queryset/aggregations.py new file mode 100644 index 000000000..654734279 --- /dev/null +++ b/ormar/queryset/aggregations.py @@ -0,0 +1,48 @@ +"""Value objects describing an aggregate function requested via ``annotate``.""" + + +class AggregateFunction: + """ + Base class for an aggregate applied to a related field or relation. + + :param field: relation name (for ``Count``) or ``relation__column`` path + :type field: str + :param distinct: whether the aggregate should be applied to distinct values + :type distinct: bool + """ + + function_name: str = "" + + def __init__(self, field: str, distinct: bool = False) -> None: + self.field = field + self.distinct = distinct + + +class Count(AggregateFunction): + """Counts related rows (or distinct related values).""" + + function_name = "count" + + +class Sum(AggregateFunction): + """Sums a related numeric column.""" + + function_name = "sum" + + +class Avg(AggregateFunction): + """Averages a related numeric column.""" + + function_name = "avg" + + +class Min(AggregateFunction): + """Minimum of a related column.""" + + function_name = "min" + + +class Max(AggregateFunction): + """Maximum of a related column.""" + + function_name = "max" diff --git a/ormar/queryset/queries/query.py b/ormar/queryset/queries/query.py index 01438d9e0..546c17a00 100644 --- a/ormar/queryset/queries/query.py +++ b/ormar/queryset/queries/query.py @@ -1,3 +1,4 @@ +import operator from typing import TYPE_CHECKING, Any, Optional, Union, cast import sqlalchemy @@ -6,7 +7,9 @@ from sqlalchemy.sql.roles import FromClauseRole import ormar # noqa I100 +from ormar.exceptions import QueryDefinitionError from ormar.models.helpers.models import group_related_list +from ormar.queryset.actions.aggregation_action import AggregationAction from ormar.queryset.actions.filter_action import FilterAction from ormar.queryset.join import SqlJoin from ormar.queryset.queries import FilterQuery, LimitQuery, OffsetQuery, OrderQuery @@ -29,6 +32,8 @@ def __init__( # noqa CFQ002 excludable: "ExcludableItems", order_bys: Optional[list["OrderAction"]], limit_raw_sql: bool, + annotations: Optional[dict] = None, + having_clauses: Optional[list] = None, ) -> None: self.query_offset = offset self.limit_count = limit_count @@ -49,6 +54,10 @@ def __init__( # noqa CFQ002 self._init_sorted_orders() self.limit_raw_sql = limit_raw_sql + self.annotations = annotations or {} + self.having_clauses = having_clauses or [] + self.aggregation_actions: list[AggregationAction] = [] + self.annotation_columns: dict = {} def _init_sorted_orders(self) -> None: """ @@ -67,13 +76,90 @@ def apply_order_bys_for_primary_model(self) -> None: # noqa: CCR001 current_table_sorted = False if self.order_columns: for clause in self.order_columns: - if clause.is_source_model_order: + if clause.field_name in self.annotations: + current_table_sorted = True + descending = clause.direction == "desc" + self.sorted_orders[clause] = self._annotations_order_text( + clause.field_name, descending + ) + elif clause.is_source_model_order: current_table_sorted = True self.sorted_orders[clause] = clause.get_text_clause() if not current_table_sorted: self._apply_default_model_sorting() + def _quote_annotation_name(self, name: str) -> str: + """ + Quotes an annotation's result column label with the dialect's own + identifier preparer (mirroring ``OrderAction.get_text_clause``), + since a hardcoded double-quote is interpreted as a string literal + rather than a column reference on dialects without ANSI_QUOTES + (e.g. MySQL). + + :param name: label of the annotation to quote + :type name: str + :return: dialect-quoted identifier + :rtype: str + """ + dialect = self.model_cls.ormar_config.database.dialect + return dialect.identifier_preparer.quote(name) + + def _annotations_order_text( + self, name: str, descending: bool + ) -> sqlalchemy.sql.expression.TextClause: + """ + Builds an ORDER BY text clause referencing an annotation's result + column by its label, since annotation labels do not correspond to + real model columns and cannot be resolved through ``OrderAction``. + + :param name: label of the annotation to order by + :type name: str + :param descending: whether to sort in descending order + :type descending: bool + :return: order by text clause referencing the annotation label + :rtype: sqlalchemy.sql.expression.TextClause + """ + quoted_name = self._quote_annotation_name(name) + direction = " desc" if descending else "" + return sqlalchemy.text(f"{quoted_name}{direction}") + + def _annotation_min_or_max_expression( + self, name: str, descending: bool + ) -> sqlalchemy.sql.ColumnElement: + """ + Builds a ``min()``/``max()``-wrapped ORDER BY expression over an + annotation's own result column, for use in the pagination subquery + built by ``_build_pagination_condition``. That subquery groups rows + by primary key only, so any other column used in its ``ORDER BY`` + (including an annotation) must be wrapped in an aggregate function; + since the annotation join is 1:1 with the parent primary key, + ``min(column) == max(column)`` is always the annotation's own value. + + This reuses ``self.annotation_columns[name]`` (the same expression + used in the main query, with e.g. ``Count``'s empty-relation + ``COALESCE(..., 0)`` already applied) rather than a bare text label: + the label also exists, unqualified, on the raw (un-coalesced) + derived-table column already present in ``self.select_from``, and + dialects disagree on how a raw ``NULL`` sorts (e.g. PostgreSQL sorts + ``NULL`` first on ``DESC`` by default), which would rank + empty-relation parents wrongly relative to the coalesced ``0`` used + everywhere else. + + :param name: label of the annotation to order by + :type name: str + :param descending: whether to sort in descending order + :type descending: bool + :return: min/max wrapped order by expression over the annotation's + own (already coalesced, where applicable) result column + :rtype: sqlalchemy.sql.ColumnElement + """ + column = self.annotation_columns[name] + wrapped = ( + sqlalchemy.func.max(column) if descending else sqlalchemy.func.min(column) + ) + return wrapped.desc() if descending else wrapped + def _apply_default_model_sorting(self) -> None: """ Applies orders_by from model OrmarConfig (if provided), if it was not provided @@ -145,6 +231,16 @@ def build_select_expression(self) -> sqlalchemy.sql.Select: self.sorted_orders, ) = sql_join.build_join() # type: ignore + for name, aggregate in self.annotations.items(): + action = AggregationAction( + name=name, aggregate=aggregate, model_cls=self.model_cls + ) + joined = action.apply_join(self.select_from, self.table) # type: ignore + self.select_from = joined # type: ignore + self.columns.append(action.result_column) # type: ignore + self.aggregation_actions.append(action) + self.annotation_columns[name] = action.result_column + if self._pagination_query_required(): limit_qry, on_clause = self._build_pagination_condition() self.select_from = sqlalchemy.sql.join( @@ -182,13 +278,17 @@ def _build_pagination_condition( pk_alias = self.model_cls.get_column_alias(self.model_cls.ormar_config.pkname) pk_aliased_name = f"{self.table.name}.{pk_alias}" qry_text = sqlalchemy.text(f"{pk_aliased_name}") - maxes = {} + maxes: dict[str, Union[TextClause, sqlalchemy.sql.ColumnElement]] = {} for order in list(self.sorted_orders.keys()): - if order is not None and order.get_field_name_text() != pk_aliased_name: + if order.field_name in self.annotations: + descending = order.direction == "desc" + maxes[order.field_name] = self._annotation_min_or_max_expression( + order.field_name, descending + ) + elif order.get_field_name_text() != pk_aliased_name: aliased_col = order.get_field_name_text() - # maxes[aliased_col] = order.get_text_clause() maxes[aliased_col] = order.get_min_or_max() - elif order.get_field_name_text() == pk_aliased_name: + else: maxes[pk_aliased_name] = order.get_text_clause() limit_qry: Select[Any] = sqlalchemy.sql.select(qry_text) @@ -197,6 +297,7 @@ def _build_pagination_condition( limit_qry = FilterQuery( filter_clauses=self.exclude_clauses, exclude=True ).apply(limit_qry) + limit_qry = self._apply_having(limit_qry) limit_qry = limit_qry.group_by(qry_text) for order_by in maxes.values(): limit_qry = limit_qry.order_by(order_by) @@ -230,12 +331,46 @@ def _apply_expression_modifiers( expr = FilterQuery(filter_clauses=self.exclude_clauses, exclude=True).apply( expr ) + expr = self._apply_having(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 + def _apply_having(self, expr: sqlalchemy.sql.Select) -> sqlalchemy.sql.Select: + """ + Applies ``having`` conditions as WHERE clauses on aggregate columns. + + Since Mode A annotations are real joined columns (not SQL + aggregates computed in the outer query), the conditions are plain + WHERE clauses rather than a SQL ``HAVING`` clause. + + :param expr: select expression before having clauses are applied + :type expr: sqlalchemy.sql.selectable.Select + :return: expression with all having clauses applied + :rtype: sqlalchemy.sql.selectable.Select + :raises QueryDefinitionError: if a having clause references a name + that was not declared through ``annotate()`` + """ + operators = { + "exact": operator.eq, + "ne": operator.ne, + "gt": operator.gt, + "gte": operator.ge, + "lt": operator.lt, + "lte": operator.le, + } + for clause in self.having_clauses: + if clause.name not in self.annotation_columns: + raise QueryDefinitionError( + f"having() references '{clause.name}' which is not an " + f"annotated aggregate; add it via annotate()." + ) + column = self.annotation_columns[clause.name] + expr = expr.where(operators[clause.op](column, clause.value)) + return expr + def _reset_query_parameters(self) -> None: """ Although it should be created each time before the call we reset the key params diff --git a/ormar/queryset/queryset.py b/ormar/queryset/queryset.py index 77b6c4f9f..f2503bc1b 100644 --- a/ormar/queryset/queryset.py +++ b/ormar/queryset/queryset.py @@ -5,6 +5,7 @@ AsyncGenerator, Generic, Iterable, + NamedTuple, Optional, Sequence, TypeVar, @@ -39,10 +40,19 @@ from ormar.models import T from ormar.models.excludable import ExcludableItems, Slot from ormar.models.ormar_config import OrmarConfig + from ormar.queryset.aggregations import AggregateFunction else: T = TypeVar("T", bound="Model") +class HavingClause(NamedTuple): + """A single parsed ``having`` condition over an annotation label.""" + + name: str + op: str + value: Any + + class QuerySet(Generic[T]): """ Main class to perform database queries, exposed on each model as objects attribute. @@ -62,6 +72,8 @@ def __init__( # noqa CFQ002 limit_raw_sql: bool = False, proxy_source_model: Optional[type["Model"]] = None, reverse_result: bool = False, + annotations: Optional[dict] = None, + having: Optional[list] = None, ) -> None: self.proxy_source_model = proxy_source_model self.model_cls = model_cls @@ -75,6 +87,8 @@ def __init__( # noqa CFQ002 self.order_bys = order_bys or [] self.limit_sql_raw = limit_raw_sql self._reverse_result = reverse_result + self._annotations = annotations or {} + self._having = having or [] @property def model_config(self) -> "OrmarConfig": @@ -113,6 +127,8 @@ def rebuild_self( # noqa: CFQ002 limit_raw_sql: Optional[bool] = None, proxy_source_model: Optional[type["Model"]] = None, reverse_result: Optional[bool] = None, + annotations: Optional[dict] = None, + having: Optional[list] = None, ) -> "QuerySet": """ Method that returns new instance of queryset based on passed params, @@ -125,6 +141,8 @@ def rebuild_self( # noqa: CFQ002 "prefetch_related": "_prefetch_related", "limit_raw_sql": "limit_sql_raw", "reverse_result": "_reverse_result", + "annotations": "_annotations", + "having": "_having", } passed_args = locals() @@ -146,6 +164,8 @@ def replace_if_none(arg_name: str) -> Any: limit_raw_sql=replace_if_none("limit_raw_sql"), proxy_source_model=replace_if_none("proxy_source_model"), reverse_result=replace_if_none("reverse_result"), + annotations=replace_if_none("annotations"), + having=replace_if_none("having"), ) async def _prefetch_related_models( @@ -292,6 +312,8 @@ def build_select_expression( order_bys=order_bys or self.order_bys, limit_raw_sql=self.limit_sql_raw, limit_count=limit if limit is not None else self.limit_count, + annotations=self._annotations, + having_clauses=self._having, ) exp = qry.build_select_expression() # print("\n", exp.compile(compile_kwargs={"literal_binds": True})) @@ -665,6 +687,51 @@ def order_by(self, columns: Union[list, str, OrderAction]) -> "QuerySet[T]": order_bys = self.order_bys + [x for x in orders_by if x not in self.order_bys] return self.rebuild_self(order_bys=order_bys) + def annotate(self, **aggregates: "AggregateFunction") -> "QuerySet[T]": + """ + Adds named aggregate columns computed per parent row via a pre-grouped + derived-table join. + + :param aggregates: mapping of result name to aggregate value object + :type aggregates: AggregateFunction + :return: queryset with the annotations applied + :rtype: QuerySet + """ + existing_names = self.model.ormar_config.model_fields + for name in aggregates: + if name in existing_names: + raise QueryDefinitionError( + f"Cannot annotate with name '{name}' - " + f"it collides with an existing field or relation on " + f"'{self.model.__name__}'. Choose a different annotation name." + ) + merged = {**self._annotations, **aggregates} + return self.rebuild_self(annotations=merged) + + def having(self, **filters: Any) -> "QuerySet[T]": + """ + Filters the queryset by annotated aggregate values. + + Suffixes mirror ``filter``: ``exact`` (default), ``gt``, ``gte``, + ``lt``, ``lte``, ``ne``. + + :param filters: mapping of ``name`` or ``name__op`` to value + :type filters: Any + :return: queryset with the having conditions applied + :rtype: QuerySet + """ + allowed = {"exact", "gt", "gte", "lt", "lte", "ne"} + clauses = list(self._having) + for key, value in filters.items(): + parts = key.split("__") + op = parts[1] if len(parts) > 1 else "exact" + if op not in allowed: + raise QueryDefinitionError( + f"Unsupported having operator '{op}'. Allowed: {sorted(allowed)}" + ) + clauses.append(HavingClause(name=parts[0], op=op, value=value)) + return self.rebuild_self(having=clauses) + async def values( self, fields: Union[list, str, set, dict, None] = None, @@ -715,8 +782,16 @@ async def values( exclude_through=exclude_through, ) column_map = alias_resolver.resolve_columns(columns_names=list(rows[0].keys())) # type: ignore + own_excludable = self._excludable.get(self.model) + annotation_names = { + name for name in self._annotations if own_excludable.is_included(name) + } result = [ - {column_map.get(k): v for k, v in dict(x).items() if k in column_map} + { + (column_map[k] if k in column_map else k): v + for k, v in dict(x).items() + if k in column_map or k in annotation_names + } for x in rows ] if _as_dict: diff --git a/tests/test_queries/test_aggregations.py b/tests/test_queries/test_aggregations.py new file mode 100644 index 000000000..e456537d2 --- /dev/null +++ b/tests/test_queries/test_aggregations.py @@ -0,0 +1,370 @@ +from typing import Optional + +import pytest +import pytest_asyncio + +import ormar +from ormar.exceptions import QueryDefinitionError +from tests.lifespan import init_tests +from tests.settings import create_config + +base_ormar_config = create_config() + + +class User(ormar.Model): + ormar_config = base_ormar_config.copy(tablename="users") + + id: int = ormar.Integer(primary_key=True) + name: str = ormar.String(max_length=100) + + +class Task(ormar.Model): + ormar_config = base_ormar_config.copy(tablename="tasks") + + id: int = ormar.Integer(primary_key=True) + user: Optional[User] = ormar.ForeignKey(User, related_name="tasks") + title: str = ormar.String(max_length=100) + price: Optional[int] = ormar.Integer(nullable=True) + + +class Tag(ormar.Model): + ormar_config = base_ormar_config.copy(tablename="tags") + + id: int = ormar.Integer(primary_key=True) + name: str = ormar.String(max_length=100) + + +class Post(ormar.Model): + ormar_config = base_ormar_config.copy(tablename="posts") + + id: int = ormar.Integer(primary_key=True) + title: str = ormar.String(max_length=100) + tags = ormar.ManyToMany(Tag, related_name="posts") + + +create_test_database = init_tests(base_ormar_config) + + +@pytest_asyncio.fixture(autouse=True, scope="function") +async def cleanup(): + yield + async with base_ormar_config.database: + await Task.objects.delete(each=True) + await User.objects.delete(each=True) + await Post.objects.delete(each=True) + await Tag.objects.delete(each=True) + + +async def seed(): + u1 = await User(name="Alice").save() + u2 = await User(name="Bob").save() + await User(name="Carol").save() # no tasks + await Task(title="a1", price=10, user=u1).save() + await Task(title="a2", price=20, user=u1).save() + await Task(title="b1", price=5, user=u2).save() + return u1, u2 + + +def test_aggregate_value_objects(): + c = ormar.Count("tasks") + assert c.function_name == "count" + assert c.field == "tasks" + assert c.distinct is False + + cd = ormar.Count("tasks", distinct=True) + assert cd.distinct is True + + assert ormar.Sum("tasks__price").function_name == "sum" + assert ormar.Avg("tasks__price").function_name == "avg" + assert ormar.Min("tasks__price").function_name == "min" + assert ormar.Max("tasks__price").function_name == "max" + assert ormar.Sum("tasks__price").field == "tasks__price" + + +@pytest.mark.asyncio +async def test_annotate_count_via_values(): + async with base_ormar_config.database: + await seed() + rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("name") + .values(["name", "task_count"]) + ) + assert rows == [ + {"name": "Alice", "task_count": 2}, + {"name": "Bob", "task_count": 1}, + {"name": "Carol", "task_count": 0}, + ] + + +@pytest.mark.asyncio +async def test_annotate_count_distinct_with_column(): + async with base_ormar_config.database: + await seed() + rows = ( + await User.objects.annotate( + task_count=ormar.Count("tasks__id", distinct=True) + ) + .order_by("name") + .values(["name", "task_count"]) + ) + assert rows == [ + {"name": "Alice", "task_count": 2}, + {"name": "Bob", "task_count": 1}, + {"name": "Carol", "task_count": 0}, + ] + + +@pytest.mark.asyncio +async def test_annotate_count_distinct_without_column(): + async with base_ormar_config.database: + await seed() + rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks", distinct=True)) + .order_by("name") + .values(["name", "task_count"]) + ) + assert rows == [ + {"name": "Alice", "task_count": 2}, + {"name": "Bob", "task_count": 1}, + {"name": "Carol", "task_count": 0}, + ] + + +@pytest.mark.asyncio +async def test_order_by_annotation(): + async with base_ormar_config.database: + await seed() + names = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("-task_count") + .values_list("name", flatten=True) + ) + assert names == ["Alice", "Bob", "Carol"] + + names_asc = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("task_count") + .values_list("name", flatten=True) + ) + assert names_asc == ["Carol", "Bob", "Alice"] + + +@pytest.mark.asyncio +async def test_having_on_annotation(): + async with base_ormar_config.database: + await seed() + rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .having(task_count__gt=1) + .values_list("name", flatten=True) + ) + assert rows == ["Alice"] + + rows_zero = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .having(task_count__gte=1) + .order_by("name") + .values_list("name", flatten=True) + ) + assert rows_zero == ["Alice", "Bob"] + + +def test_having_with_unsupported_operator(): + with pytest.raises(QueryDefinitionError): + User.objects.annotate(task_count=ormar.Count("tasks")).having( + task_count__contains=1 + ) + + +@pytest.mark.asyncio +async def test_sum_avg_min_max_annotations(): + async with base_ormar_config.database: + await seed() + rows = ( + await User.objects.annotate( + total=ormar.Sum("tasks__price"), + avg_price=ormar.Avg("tasks__price"), + cheapest=ormar.Min("tasks__price"), + dearest=ormar.Max("tasks__price"), + ) + .filter(name="Alice") + .values(["name", "total", "cheapest", "dearest"]) + ) + assert rows == [{"name": "Alice", "total": 30, "cheapest": 10, "dearest": 20}] + + # parent with no children => NULL (not 0) for non-count aggregates + carol = ( + await User.objects.annotate(total=ormar.Sum("tasks__price")) + .filter(name="Carol") + .values(["name", "total"]) + ) + assert carol == [{"name": "Carol", "total": None}] + + +@pytest.mark.asyncio +async def test_count_m2m_annotation(): + async with base_ormar_config.database: + t1 = await Tag(name="t1").save() + t2 = await Tag(name="t2").save() + p1 = await Post(title="p1").save() + p2 = await Post(title="p2").save() + await p1.tags.add(t1) + await p1.tags.add(t2) + await p2.tags.add(t1) + rows = ( + await Post.objects.annotate(tag_count=ormar.Count("tags")) + .order_by("title") + .values(["title", "tag_count"]) + ) + assert rows == [ + {"title": "p1", "tag_count": 2}, + {"title": "p2", "tag_count": 1}, + ] + + +@pytest.mark.asyncio +async def test_qualified_m2m_aggregate_raises(): + async with base_ormar_config.database: + with pytest.raises(QueryDefinitionError): + await Post.objects.annotate(x=ormar.Sum("tags__id")).values(["title", "x"]) + + +@pytest.mark.asyncio +async def test_annotation_with_select_related_and_limit(): + async with base_ormar_config.database: + await seed() + # select_related hydrates tasks; annotation count stays correct and 1 row/user + users = ( + await User.objects.select_related("tasks") + .annotate(task_count=ormar.Count("tasks")) + .order_by("-task_count") + .limit(2) + .all() + ) + assert [u.name for u in users] == ["Alice", "Bob"] + # ascending order differs from insertion/pk order (Alice, Bob, Carol), + # so this discriminates a genuinely applied sort from an inert one + users_asc = ( + await User.objects.select_related("tasks") + .annotate(task_count=ormar.Count("tasks")) + .order_by("task_count") + .limit(2) + .all() + ) + assert [u.name for u in users_asc] == ["Carol", "Bob"] + # offset skips the top-ranked row (Alice); an inert ORDER BY combined + # with offset would instead skip Bob and keep Alice in the results + users_offset = ( + await User.objects.select_related("tasks") + .annotate(task_count=ormar.Count("tasks")) + .order_by("-task_count") + .offset(1) + .limit(2) + .all() + ) + assert [u.name for u in users_offset] == ["Bob", "Carol"] + # values() with limit + annotation order + rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("-task_count") + .limit(1) + .values(["name", "task_count"]) + ) + assert rows == [{"name": "Alice", "task_count": 2}] + + +@pytest.mark.asyncio +async def test_annotate_values_without_field_subset(): + async with base_ormar_config.database: + await seed() + rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .order_by("name") + .values() + ) + assert len(rows) == 3 + by_name = {row["name"]: row for row in rows} + assert by_name["Alice"]["task_count"] == 2 + assert by_name["Bob"]["task_count"] == 1 + assert by_name["Carol"]["task_count"] == 0 + # regular (non-annotation) columns are still present too + assert "id" in by_name["Alice"] + + +@pytest.mark.asyncio +async def test_having_with_select_related_and_limit(): + async with base_ormar_config.database: + await seed() # Alice (pk 1, 2 tasks), Bob (pk 2, 1 task), Carol (pk 3, 0 tasks) + dave = await User(name="Dave").save() # pk 4 + for i in range(3): + await Task(title=f"d{i}", price=1, user=dave).save() + # pk order is Alice(2 tasks), Bob(1 task), Carol(0 tasks), Dave(3 tasks). + # Only Alice and Dave satisfy `having(task_count__gt=1)`. A pagination + # subquery that picks the first-by-pk 2 parents *before* having is + # applied would pick Alice and Bob, then the outer filter would drop + # Bob (1 task), silently returning only [Alice] instead of the correct + # [Alice, Dave] - this is discriminating against the Fix-1 regression. + users = ( + await User.objects.select_related("tasks") + .annotate(task_count=ormar.Count("tasks")) + .having(task_count__gt=1) + .limit(2) + .all() + ) + assert sorted(u.name for u in users) == ["Alice", "Dave"] + assert len(users) == 2 + + +def test_annotate_name_collision_raises(): + with pytest.raises(QueryDefinitionError): + User.objects.annotate(name=ormar.Count("tasks")) + + +@pytest.mark.asyncio +async def test_having_on_unannotated_name_raises(): + async with base_ormar_config.database: + await seed() + with pytest.raises(QueryDefinitionError): + await User.objects.having(missing__gt=1).all() + + +@pytest.mark.asyncio +async def test_annotation_independent_of_outer_relation_filter(): + async with base_ormar_config.database: + u1, _ = await seed() + await Task(title="a3", price=50, user=u1).save() # Alice now has 3 tasks + rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .filter(tasks__price__gt=15) + .order_by("name") + .values(["name", "task_count"]) + ) + # The outer `.filter(tasks__price__gt=15)` joins `tasks` again (aliased) + # purely to filter rows; it matches 2 of Alice's tasks (price 20 and + # 50), so plain SQL join semantics duplicate her row in the flat + # `.values()` result (pre-existing ormar behaviour, unrelated to + # annotations - see the `distinct` flag on `QuerySet.count()`). + # What this test locks down is that `task_count` is unaffected by + # that outer filter: it is computed by a separate derived-table + # subquery grouped over the raw child table, so it is `3` (ALL of + # Alice's tasks) on every duplicated row, not `2` (only the tasks + # matching the outer filter). + assert rows == [ + {"name": "Alice", "task_count": 3}, + {"name": "Alice", "task_count": 3}, + ] + + # Same independence property, without the duplication confound: + # a scalar-column filter (`name="Alice"`) matches Alice's row once, + # so no join-induced duplication occurs, yet `task_count` is still + # `3` (ALL of Alice's tasks). A to-many-relation filter duplicates + # rows (asserted above); a scalar filter does not (asserted below) - + # in both cases `task_count` reflects every task, independent of + # the outer filter. + scalar_filtered_rows = ( + await User.objects.annotate(task_count=ormar.Count("tasks")) + .filter(name="Alice") + .values(["name", "task_count"]) + ) + assert scalar_filtered_rows == [{"name": "Alice", "task_count": 3}]