From d53516d4689b2cf451148896b592d7f3295b251c Mon Sep 17 00:00:00 2001 From: long2ice Date: Thu, 30 Apr 2020 18:49:42 +0800 Subject: [PATCH] support group_by with join (#374) --- CHANGELOG.rst | 1 + examples/group_by.py | 9 +++++++ tests/test_group_by.py | 55 ++++++++++++++++++++++++++++++++++++++++++ tortoise/queryset.py | 20 ++++++++++++--- 4 files changed, 81 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 574347707..979ab48ef 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,7 @@ Changelog ------- - Fix bad import of ``basestring`` - Better handling of NULL characters in strings. Fixes SQLite, raises better error for PostgreSQL. +- Support ``.group_by()`` with join now 0.16.9 ------ diff --git a/examples/group_by.py b/examples/group_by.py index de821319a..d1c38e01f 100644 --- a/examples/group_by.py +++ b/examples/group_by.py @@ -71,6 +71,15 @@ async def run(): print(ret) # >>> [(1, 10), (2, 5)] + # group by with join + ret = ( + await Book.annotate(count=Count("id")) + .group_by("author__name") + .values("author__name", "count") + ) + print(ret) + # >>> [{"author__name": "author1", "count": 10}, {"author__name": "author2", "count": 5}] + if __name__ == "__main__": run_async(run()) diff --git a/tests/test_group_by.py b/tests/test_group_by.py index 1c1953412..d6129bf57 100644 --- a/tests/test_group_by.py +++ b/tests/test_group_by.py @@ -27,6 +27,16 @@ async def test_count_group_by(self): elif author_id == self.a2.pk: self.assertEqual(count, 5) + async def test_count_group_by_with_join(self): + ret = ( + await Book.annotate(count=Count("id")) + .group_by("author__name") + .values("author__name", "count") + ) + self.assertEqual( + ret, [{"author__name": "author1", "count": 10}, {"author__name": "author2", "count": 5}] + ) + async def test_count_filter_group_by(self): ret = ( await Book.annotate(count=Count("id")) @@ -49,6 +59,17 @@ async def test_sum_group_by(self): elif author_id == self.a2.pk: self.assertEqual(sum_, 10.0) + async def test_sum_group_by_with_join(self): + ret = ( + await Book.annotate(sum=Sum("rating")) + .group_by("author__name") + .values("author__name", "sum") + ) + self.assertEqual( + ret, + [{"author__name": "author1", "sum": 45.0}, {"author__name": "author2", "sum": 10.0}], + ) + async def test_sum_filter_group_by(self): ret = ( await Book.annotate(sum=Sum("rating")) @@ -72,6 +93,16 @@ async def test_avg_group_by(self): elif author_id == self.a2.pk: self.assertEqual(avg, 2.0) + async def test_avg_group_by_with_join(self): + ret = ( + await Book.annotate(avg=Avg("rating")) + .group_by("author__name") + .values("author__name", "avg") + ) + self.assertEqual( + ret, [{"author__name": "author1", "avg": 4.5}, {"author__name": "author2", "avg": 2}] + ) + async def test_avg_filter_group_by(self): ret = ( await Book.annotate(avg=Avg("rating")) @@ -97,6 +128,14 @@ async def test_count_values_list_group_by(self): elif author_id == self.a2.pk: self.assertEqual(count, 5) + async def test_count_values_list_group_by_with_join(self): + ret = ( + await Book.annotate(count=Count("id")) + .group_by("author__name") + .values_list("author__name", "count") + ) + self.assertEqual(ret, [("author1", 10), ("author2", 5)]) + async def test_count_values_list_filter_group_by(self): ret = ( await Book.annotate(count=Count("id")) @@ -121,6 +160,14 @@ async def test_sum_values_list_group_by(self): elif author_id == self.a2.pk: self.assertEqual(sum_, 10.0) + async def test_sum_values_list_group_by_with_join(self): + ret = ( + await Book.annotate(sum=Sum("rating")) + .group_by("author__name") + .values_list("author__name", "sum") + ) + self.assertEqual(ret, [("author1", 45.0), ("author2", 10.0)]) + async def test_sum_values_list_filter_group_by(self): ret = ( await Book.annotate(sum=Sum("rating")) @@ -146,6 +193,14 @@ async def test_avg_values_list_group_by(self): elif author_id == self.a2.pk: self.assertEqual(avg, 2.0) + async def test_avg_values_list_group_by_with_join(self): + ret = ( + await Book.annotate(avg=Avg("rating")) + .group_by("author__name") + .values_list("author__name", "avg") + ) + self.assertEqual(ret, [("author1", 4.5), ("author2", 2.0)]) + async def test_avg_values_list_filter_group_by(self): ret = ( await Book.annotate(avg=Avg("rating")) diff --git a/tortoise/queryset.py b/tortoise/queryset.py index d8bfc52f7..60eefa503 100644 --- a/tortoise/queryset.py +++ b/tortoise/queryset.py @@ -948,6 +948,20 @@ def resolve_to_python_value(self, model: Type[MODEL], field: str) -> Callable: raise FieldError(f'Unknown field "{field}" for model "{model}"') + def _resolve_group_bys(self, *field_names: str): + group_bys = [] + for field_name in field_names: + field_split = field_name.split("__") + related_table, related_db_field = self._join_table_with_forwarded_fields( + model=self.model, + table=self.model._meta.basetable, + field=field_split[0], + forwarded_fields="__".join(field_split[1:]) if len(field_split) > 1 else "", + ) + field = related_table[related_db_field].as_(field_name) + group_bys.append(field) + return group_bys + class ValuesListQuery(FieldSelectQuery): __slots__ = ( @@ -1017,8 +1031,7 @@ def _make_query(self) -> None: if self.distinct: self.query._distinct = True if self.group_bys: - self.query._groupbys = [] - self.query = self.query.groupby(*self.group_bys) + self.query._groupbys = self._resolve_group_bys(*self.group_bys) def __await__(self) -> Generator[Any, None, List[Any]]: if self._db is None: @@ -1104,8 +1117,7 @@ def _make_query(self) -> None: if self.distinct: self.query._distinct = True if self.group_bys: - self.query._groupbys = [] - self.query = self.query.groupby(*self.group_bys) + self.query._groupbys = self._resolve_group_bys(*self.group_bys) def __await__(self) -> Generator[Any, None, List[dict]]: if self._db is None: