fix: fix wrong total_count when using distinct on m2m/o2m relationships

Fix #792

Author: bellini666

strawberry_django/pagination.py

@@ -249,8 +249,8 @@ def apply_window_pagination(
 def remove_window_pagination(queryset: _QS) -> _QS:
     """Remove pagination window functions from a queryset.
 
-    Utility function to remove the pagination `WHERE` clause added by
-    the `apply_window_pagination` function.
+    Utility function to remove the pagination `WHERE` clause and annotations
+    added by the `apply_window_pagination` function.
 
     Args:
     ----
@@ -263,6 +263,11 @@ def remove_window_pagination(queryset: _QS) -> _QS:
         for child in queryset.query.where.children
         if (not hasattr(child, "lhs") or not isinstance(child.lhs, _PaginationWindow))
     ]
+    queryset.query.annotations = {  # type: ignore
+        key: value
+        for key, value in queryset.query.annotations.items()
+        if not isinstance(value, _PaginationWindow)
+    }
     return queryset
 
 
@@ -278,6 +283,13 @@ def get_total_count(queryset: QuerySet) -> int:
         results = queryset._result_cache  # type: ignore
 
         if results:
+            # If the queryset has DISTINCT enabled, the _strawberry_total_count
+            # annotation won't be accurate because window functions are evaluated
+            # before DISTINCT in SQL. Fall back to queryset.count() instead.
+            if queryset.query.distinct:
+                queryset = remove_window_pagination(queryset)
+                return queryset.count()
+
             try:
                 return results[0]._strawberry_total_count
             except AttributeError:

tests/relay/test_cursor_pagination.py

@@ -1419,6 +1419,152 @@ def test_invalid_cursor(cursor, test_objects):
     assert result.errors[0].message == "Invalid cursor"
 
 
+@pytest.mark.django_db(transaction=True)
+def test_connection_distinct_with_m2m_filters():
+    from tests import models
+
+    @strawberry_django.filter_type(models.FruitType, lookups=True)
+    class FruitTypeFilter:
+        name: strawberry.auto
+
+    @strawberry_django.filter_type(models.Fruit, lookups=True)
+    class FruitFilter:
+        name: strawberry.auto
+        types: FruitTypeFilter | None
+        DISTINCT: bool | None
+
+    @strawberry_django.type(models.FruitType)
+    class FruitTypeGQL(Node):
+        name: strawberry.auto
+
+    @strawberry_django.type(models.Fruit, filters=FruitFilter)
+    class FruitGQL(Node):
+        name: strawberry.auto
+        types: list[FruitTypeGQL]
+
+    @strawberry.type
+    class FruitQuery:
+        fruits: DjangoCursorConnection[FruitGQL] = strawberry_django.connection()
+
+    fruit_schema = strawberry.Schema(
+        query=FruitQuery, extensions=[DjangoOptimizerExtension()]
+    )
+
+    ft1 = models.FruitType.objects.create(name="tropical")
+    ft2 = models.FruitType.objects.create(name="citrus")
+
+    banana = models.Fruit.objects.create(name="banana")
+    banana.types.add(ft1, ft2)
+
+    apple = models.Fruit.objects.create(name="apple")
+    apple.types.add(ft1)
+
+    orange = models.Fruit.objects.create(name="orange")
+    orange.types.add(ft2)
+
+    query = """
+    query TestQuery {
+        fruits(filters: {
+            types: { name: { inList: ["tropical", "citrus"] } }
+            DISTINCT: true
+        }) {
+            edges {
+                node {
+                    id
+                    name
+                }
+            }
+            totalCount
+        }
+    }
+    """
+    result = fruit_schema.execute_sync(query)
+    assert not result.errors
+    assert result.data
+
+    edges = result.data["fruits"]["edges"]
+    total_count = result.data["fruits"]["totalCount"]
+
+    assert len(edges) == 3
+    assert total_count == 3
+
+    names = {edge["node"]["name"] for edge in edges}
+    assert names == {"banana", "apple", "orange"}
+
+
+@pytest.mark.django_db(transaction=True)
+def test_connection_distinct_with_pagination():
+    from tests import models
+
+    @strawberry_django.filter_type(models.FruitType, lookups=True)
+    class FruitTypeFilter:
+        name: strawberry.auto
+
+    @strawberry_django.filter_type(models.Fruit, lookups=True)
+    class FruitFilter:
+        name: strawberry.auto
+        types: FruitTypeFilter | None
+        DISTINCT: bool | None
+
+    @strawberry_django.type(models.FruitType)
+    class FruitTypeGQL(Node):
+        name: strawberry.auto
+
+    @strawberry_django.type(models.Fruit, filters=FruitFilter)
+    class FruitGQL(Node):
+        name: strawberry.auto
+        types: list[FruitTypeGQL]
+
+    @strawberry.type
+    class FruitQuery:
+        fruits: DjangoCursorConnection[FruitGQL] = strawberry_django.connection()
+
+    fruit_schema = strawberry.Schema(
+        query=FruitQuery, extensions=[DjangoOptimizerExtension()]
+    )
+
+    ft1 = models.FruitType.objects.create(name="tropical")
+    ft2 = models.FruitType.objects.create(name="citrus")
+
+    for i in range(5):
+        fruit = models.Fruit.objects.create(name=f"fruit_{i}")
+        fruit.types.add(ft1, ft2)
+
+    query = """
+    query TestQuery {
+        fruits(
+            first: 2
+            filters: {
+                types: { name: { inList: ["tropical", "citrus"] } }
+                DISTINCT: true
+            }
+        ) {
+            edges {
+                node {
+                    id
+                    name
+                }
+            }
+            totalCount
+            pageInfo {
+                hasNextPage
+            }
+        }
+    }
+    """
+    result = fruit_schema.execute_sync(query)
+    assert not result.errors
+    assert result.data
+
+    edges = result.data["fruits"]["edges"]
+    total_count = result.data["fruits"]["totalCount"]
+    has_next = result.data["fruits"]["pageInfo"]["hasNextPage"]
+
+    assert len(edges) == 2
+    assert total_count == 5
+    assert has_next is True
+
+
 @pytest.mark.django_db(transaction=True)
 @pytest.mark.parametrize(
     ("first", "last", "error_message"),