Merge pull request #589 from erichaydel/588-fix-queryset-count-bug

Fix collerek/ormar#588 Bug in queryset count() method

README.md

@@ -601,7 +601,7 @@ metadata.drop_all(engine)
 *  `prefetch_related(related: Union[List, str]) -> QuerySet`
 *  `limit(limit_count: int) -> QuerySet`
 *  `offset(offset: int) -> QuerySet`
-*  `count() -> int`
+*  `count(distinct: bool = True) -> int`
 *  `exists() -> bool`
 *  `max(columns: List[str]) -> Any`
 *  `min(columns: List[str]) -> Any`

docs/api/query-set/query-set.md

@@ -499,11 +499,19 @@ Returns a bool value to confirm if there are rows matching the given criteria
 #### count
 
 ```python
- | async count() -> int
+ | async count(distinct: bool = True) -> int
 ```
 
 Returns number of rows matching the given criteria
 (applied with `filter` and `exclude` if set before).
+If `distinct` is `True` (the default), this will return the number of primary rows selected. If `False`,
+the count will be the total number of rows returned
+(including extra rows for `one-to-many` or `many-to-many` left `select_related` table joins).
+`False` is the legacy (buggy) behavior for workflows that depend on it.
+
+**Arguments**:
+
+- `distinct` (`bool`): flag if the primary table rows should be distinct or not
 
 **Returns**:
 
@@ -865,4 +873,3 @@ Bulk operations do not send signals.
 
 - `objects` (`List[Model]`): list of ormar models
 - `columns` (`List[str]`): list of columns to update
-

docs/api/relations/queryset-proxy.md

@@ -150,14 +150,22 @@ Actual call delegated to QuerySet.
 #### count
 
 ```python
- | async count() -> int
+ | async count(distinct: bool = True) -> int
 ```
 
 Returns number of rows matching the given criteria
 (applied with `filter` and `exclude` if set before).
+If `distinct` is `True` (the default), this will return the number of primary rows selected. If `False`,
+the count will be the total number of rows returned
+(including extra rows for `one-to-many` or `many-to-many` left `select_related` table joins).
+`False` is the legacy (buggy) behavior for workflows that depend on it.
 
 Actual call delegated to QuerySet.
 
+**Arguments**:
+
+- `distinct` (`bool`): flag if the primary table rows should be distinct or not
+
 **Returns**:
 
 `int`: number of rows
@@ -773,4 +781,3 @@ Actual call delegated to QuerySet.
 **Returns**:
 
 `QuerysetProxy`: QuerysetProxy
-

docs/index.md

@@ -610,7 +610,7 @@ metadata.drop_all(engine)
 *  `prefetch_related(related: Union[List, str]) -> QuerySet`
 *  `limit(limit_count: int) -> QuerySet`
 *  `offset(offset: int) -> QuerySet`
-*  `count() -> int`
+*  `count(distinct: bool = True) -> int`
 *  `exists() -> bool`
 *  `max(columns: List[str]) -> Any`
 *  `min(columns: List[str]) -> Any`

docs/queries/aggregations.md

@@ -3,7 +3,7 @@
 Currently 6 aggregation functions are supported.
 
 
-* `count() -> int`
+* `count(distinct: bool = True) -> int`
 * `exists() -> bool`
 * `sum(columns) -> Any`
 * `avg(columns) -> Any`
@@ -12,7 +12,7 @@ Currently 6 aggregation functions are supported.
 
 
 * `QuerysetProxy`
-    * `QuerysetProxy.count()` method
+    * `QuerysetProxy.count(distinct=True)` method
     * `QuerysetProxy.exists()` method
     * `QuerysetProxy.sum(columns)` method
     * `QuerysetProxy.avg(columns)` method
@@ -22,9 +22,13 @@ Currently 6 aggregation functions are supported.
 
 ## count
 
-`count() -> int`
+`count(distinct: bool = True) -> int`
 
-Returns number of rows matching the given criteria (i.e. applied with `filter` and `exclude`)
+Returns number of rows matching the given criteria (i.e. applied with `filter` and `exclude`).
+If `distinct` is `True` (the default), this will return the number of primary rows selected. If `False`,
+the count will be the total number of rows returned
+(including extra rows for `one-to-many` or `many-to-many` left `select_related` table joins).
+`False` is the legacy (buggy) behavior for workflows that depend on it.
 
 ```python
 class Book(ormar.Model):
@@ -322,4 +326,3 @@ objects from other side of the relation.
 
 !!!tip
     To read more about `QuerysetProxy` visit [querysetproxy][querysetproxy] section
-

docs/queries/index.md

@@ -188,12 +188,12 @@ Instead of ormar models return raw data in form list of dictionaries or tuples.
 
 ### [Aggregated functions](./aggregations.md)
 
-* `count() -> int`
+* `count(distinct: bool = True) -> int`
 * `exists() -> bool`
 
 
 * `QuerysetProxy`
-    * `QuerysetProxy.count()` method
+    * `QuerysetProxy.count(distinct=True)` method
     * `QuerysetProxy.exists()` method
 
 !!!tip

docs/relations/queryset-proxy.md

@@ -274,7 +274,7 @@ With exclude_fields() you can select subset of model columns that will be exclud
 
 ### count
 
-`count() -> int`
+`count(distinct: bool = True) -> int`
 
 Returns number of rows matching the given criteria (i.e. applied with filter and exclude)
 

ormar/protocols/queryset_protocol.py

@@ -26,7 +26,7 @@ class QuerySetProtocol(Protocol):  # pragma: nocover
     async def exists(self) -> bool:
         ...
 
-    async def count(self) -> int:
+    async def count(self, distinct: bool = True) -> int:
         ...
 
     async def clear(self) -> int:

ormar/queryset/queryset.py

@@ -678,16 +678,25 @@ class QuerySet(Generic[T]):
         expr = sqlalchemy.exists(expr).select()
         return await self.database.fetch_val(expr)
 
-    async def count(self) -> int:
+    async def count(self, distinct: bool = True) -> int:
         """
         Returns number of rows matching the given criteria
         (applied with `filter` and `exclude` if set before).
+        If `distinct` is `True` (the default), this will return the number of primary rows selected. If `False`,
+        the count will be the total number of rows returned
+        (including extra rows for `one-to-many` or `many-to-many` left `select_related` table joins).
+        `False` is the legacy (buggy) behavior for workflows that depend on it.
+
+        :param distinct: flag if the primary table rows should be distinct or not
 
         :return: number of rows
         :rtype: int
         """
         expr = self.build_select_expression().alias("subquery_for_count")
         expr = sqlalchemy.func.count().select().select_from(expr)
+        if distinct:
+            expr_distinct = expr.group_by(self.model_meta.pkname).alias("subquery_for_group")
+            expr = sqlalchemy.func.count().select().select_from(expr_distinct)
         return await self.database.fetch_val(expr)
 
     async def _query_aggr_function(self, func_name: str, columns: List) -> Any:

ormar/relations/querysetproxy.py

@@ -193,17 +193,22 @@ class QuerysetProxy(Generic[T]):
         """
         return await self.queryset.exists()
 
-    async def count(self) -> int:
+    async def count(self, distinct: bool = True) -> int:
         """
         Returns number of rows matching the given criteria
         (applied with `filter` and `exclude` if set before).
+        If `distinct` is `True` (the default), this will return the number of primary rows selected. If `False`,
+        the count will be the total number of rows returned
+        (including extra rows for `one-to-many` or `many-to-many` left `select_related` table joins).
+        `False` is the legacy (buggy) behavior for workflows that depend on it.
 
         Actual call delegated to QuerySet.
 
+        :param distinct: flag if the primary table rows should be distinct or not
         :return: number of rows
         :rtype: int
         """
-        return await self.queryset.count()
+        return await self.queryset.count(distinct=distinct)
 
     async def max(self, columns: Union[str, List[str]]) -> Any:  # noqa: A003
         """

tests/test_queries/test_aggr_functions.py

@@ -175,3 +175,15 @@ async def test_queryset_method():
         assert await author.books.max(["year", "title"]) == dict(
             year=1930, title="Book 3"
         )
+
+@pytest.mark.asyncio
+async def test_count_method():
+    async with database:
+        await sample_data()
+
+        count =  await Author.objects.select_related("books").count()
+        assert count == 1
+
+        # The legacy functionality
+        count =  await Author.objects.select_related("books").count(distinct=False)
+        assert count == 3