From b0a3aaab6c2a698aca7c4e2f2829d8518293a5fb Mon Sep 17 00:00:00 2001 From: Ivan Larin Date: Fri, 10 Jan 2025 18:31:24 +0300 Subject: [PATCH] =?UTF-8?q?feat(queries):=20=D1=81ount=20query?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #294 --- docs/queries.md | 43 +++++++++++++++++++++- edgeql_qb/queries.py | 19 ++++++++++ edgeql_qb/render/count.py | 34 +++++++++++++++++ tests/test_renderer/test_count_renderer.py | 26 +++++++++++++ 4 files changed, 121 insertions(+), 1 deletion(-) create mode 100644 edgeql_qb/render/count.py create mode 100644 tests/test_renderer/test_count_renderer.py diff --git a/docs/queries.md b/docs/queries.md index 4d4e512..4e02fcf 100644 --- a/docs/queries.md +++ b/docs/queries.md @@ -173,9 +173,50 @@ as context's variable `limit_N` and will pass to EdgeDB dynamically, which will Similar to how it works in SQLAlchemy's `unsafe_text` wrapper make this expression hardcoded into final query as is, without dynamic contexts. -Please note that `offset` by design producing not optional execution plan +Please note that `offset` by design producing not optional execution plan, and you have to avoid to use this keyword and method as far as you can. +### Counting + +The count method allows you to calculate the total number of objects returned by a query. +It supports the application of filters and other query modifiers such as where. + +#### Count Without Filters + +To count all objects of a model: + +```python +Movie.count.build() +``` + +
+ generated query + +``` +select count(Movie) +``` +
+ +#### Count with Filters + +You can apply filters using the where method to count a subset of objects. For example: + +```python +Movie.count.where(Movie.c.year <= int16(2000)).build() +``` + +
+ generated query + +``` +select count((select Movie filter .year <= $filter_0)) +{'filter_0': 2000} +``` +
+ +Please note that `select count(Movie) filter .year < 2000` +and `select count((select Movie filter .year < 2000))` have different semantics in EdgeDB. +This library implements only the second one. ## Insert diff --git a/edgeql_qb/queries.py b/edgeql_qb/queries.py index c919b72..c58075d 100644 --- a/edgeql_qb/queries.py +++ b/edgeql_qb/queries.py @@ -16,6 +16,7 @@ from edgeql_qb.func import FuncInvocation from edgeql_qb.operators import BinaryOp, SortedExpression, UnaryOp from edgeql_qb.render.condition import render_conditions +from edgeql_qb.render.count import render_count, render_count_inner from edgeql_qb.render.delete import render_delete from edgeql_qb.render.group import ( render_group, @@ -47,6 +48,10 @@ def select(self, *selectables: SelectExpressions, **kwargs: SubQuery) -> 'Select _select=(*select_args, *select_kwargs), ) + @property + def count(self): + return CountQuery(_model=self) + def group(self, *selectables: SelectExpressions) -> 'GroupQuery': return GroupQuery( _model=self, @@ -127,6 +132,20 @@ def build(self, generator: Iterator[int] | None = None) -> RenderedQuery: ) +@dataclass(slots=True, frozen=True) +class CountQuery: + _model: EdgeDBModel + _filters: tuple[Expression, ...] = field(default_factory=tuple) + + def where(self, compared: BinaryOp | UnaryOp | FuncInvocation) -> 'SelectQuery': + return replace(self, _filters=(*self._filters, Expression(compared))) + + def build(self, generator: Iterator[int] | None = None) -> RenderedQuery: + gen = generator or count() + inner = render_count_inner(self._model.name, self._filters, gen) + return render_count(inner) + + @dataclass(slots=True, frozen=True) class GroupQuery: _model: EdgeDBModel diff --git a/edgeql_qb/render/count.py b/edgeql_qb/render/count.py new file mode 100644 index 0000000..361f048 --- /dev/null +++ b/edgeql_qb/render/count.py @@ -0,0 +1,34 @@ +from collections.abc import Iterator + +from mypy.nodes import Expression + +from edgeql_qb.render.condition import render_conditions +from edgeql_qb.render.select import render_select +from edgeql_qb.render.tools import combine_many_renderers +from edgeql_qb.render.types import RenderedQuery + + +def render_count(inner: RenderedQuery) -> RenderedQuery: + return combine_many_renderers( + RenderedQuery('select count('), + inner, + RenderedQuery(')'), + ) + +def render_count_inner(model_name, filters: tuple[Expression, ...], gen: Iterator[int]): + match filters: + case (): + return RenderedQuery(model_name) + case conditions: + rendered_select = render_select( + model_name, + select=(), + generator=gen, + ) + rendered_filters = render_conditions(conditions, gen) + return combine_many_renderers( + RenderedQuery('('), + rendered_select, + rendered_filters, + RenderedQuery(')'), + ) diff --git a/tests/test_renderer/test_count_renderer.py b/tests/test_renderer/test_count_renderer.py new file mode 100644 index 0000000..9ceb81d --- /dev/null +++ b/tests/test_renderer/test_count_renderer.py @@ -0,0 +1,26 @@ +from multiprocessing.connection import Client + +from edgeql_qb import EdgeDBModel +from edgeql_qb.frozendict import FrozenDict +from edgeql_qb.types import int16 + +A = EdgeDBModel('A') + + +def test_count_wo_filters() -> None: + rendered = A.count.build() + assert rendered.query == 'select count(A)' + + +def test_count_filter(client: Client) -> None: + insert1 = A.insert.values(p_int16=int16(1)).build() + insert2 = A.insert.values(p_int16=int16(2)).build() + insert3 = A.insert.values(p_int16=int16(3)).build() + client.query(insert1.query, **insert1.context) + client.query(insert2.query, **insert2.context) + client.query(insert3.query, **insert3.context) + rendered = A.count.where(A.c.p_int16 <= int16(2)).build() + assert rendered.query == 'select count((select A filter .p_int16 <= $filter_0))' + assert rendered.context == FrozenDict(filter_0=2) + result = client.query(rendered.query, **rendered.context) + assert result == [2]