Relay pagination optimizations (#777)

Author: Kitefiko

strawberry_django/optimizer.py

@@ -37,7 +37,7 @@
 from graphql.execution.collect_fields import collect_sub_fields
 from graphql.language.ast import OperationType
 from graphql.type.definition import GraphQLResolveInfo, get_named_type
-from strawberry import relay
+from strawberry import UNSET, relay
 from strawberry.extensions import SchemaExtension
 from strawberry.relay.utils import SliceMetadata
 from strawberry.schema.schema import Schema
@@ -82,10 +82,13 @@
     from collections.abc import Generator
 
     from django.contrib.contenttypes.fields import GenericRelation
+    from strawberry.relay import Edge
     from strawberry.types.execution import ExecutionContext
     from strawberry.types.field import StrawberryField
     from strawberry.utils.await_maybe import AwaitableOrValue
 
+    from strawberry_django.pagination import OffsetPaginationInfo
+
 
 __all__ = [
     "DjangoOptimizerExtension",
@@ -582,19 +585,33 @@ def _optimize_prefetch_queryset(
                 connection_type is relay.ListConnection
                 or connection_type is DjangoListConnection
             ):
+                field_def_ = connection_type_def.get_field("edges")
+                assert field_def_
+                field_ = field_def_.resolve_type(type_definition=connection_type_def)
+                field_ = unwrap_type(field_)
+                edge_class = cast("Edge", field_)
+
                 slice_metadata = SliceMetadata.from_arguments(
                     Info(_raw_info=info, _field=field),
                     first=field_kwargs.get("first"),
                     last=field_kwargs.get("last"),
                     before=field_kwargs.get("before"),
                     after=field_kwargs.get("after"),
+                    max_results=connection_extension.max_results,
+                    prefix=edge_class.CURSOR_PREFIX,
                 )
+                mark_reversed = slice_metadata.expected is None
                 qs = apply_window_pagination(
                     qs,
                     related_field_id=related_field_id,
                     offset=slice_metadata.start,
-                    limit=slice_metadata.end - slice_metadata.start,
+                    limit=(
+                        field_kwargs.get("last", UNSET)
+                        if mark_reversed
+                        else slice_metadata.end - slice_metadata.start
+                    ),
                     max_results=connection_extension.max_results,
+                    reverse=mark_reversed,
                 )
             elif connection_type is DjangoCursorConnection:
                 qs, _ = apply_cursor_pagination(
@@ -611,7 +628,7 @@ def _optimize_prefetch_queryset(
                 mark_optimized = False
 
         if isinstance(field.type, type) and issubclass(field.type, OffsetPaginated):
-            pagination = field_kwargs.get("pagination")
+            pagination: OffsetPaginationInfo | None = field_kwargs.get("pagination")
             qs = apply_window_pagination(
                 qs,
                 related_field_id=related_field_id,

strawberry_django/pagination.py

@@ -172,6 +172,7 @@ def apply_window_pagination(
     offset: int = 0,
     limit: Optional[int] = UNSET,
     max_results: Optional[int] = None,
+    reverse: bool = False,
 ) -> _QS:
     """Apply pagination using window functions.
 
@@ -186,8 +187,17 @@ def apply_window_pagination(
         related_field_id: The related field id to apply pagination to.
         offset: The offset to start the pagination from.
         limit: The limit of items to return.
+        reverse: The need to reverse queryset ordering for backwards relay pagination
 
     """
+    if limit is UNSET:
+        settings = strawberry_django_settings()
+        limit = (
+            max_results
+            if max_results is not None
+            else settings["PAGINATION_DEFAULT_LIMIT"]
+        )
+
     order_by = [
         expr
         for expr, _ in queryset.query.get_compiler(
@@ -210,13 +220,23 @@ def apply_window_pagination(
     if offset:
         queryset = queryset.filter(_strawberry_row_number__gt=offset)
 
-    if limit is UNSET:
-        settings = strawberry_django_settings()
-        limit = (
-            max_results
-            if max_results is not None
-            else settings["PAGINATION_DEFAULT_LIMIT"]
+    if reverse:
+        order_by_reverse = [
+            expr
+            for expr, _ in queryset.reverse()
+            .query.get_compiler(
+                using=queryset._db or DEFAULT_DB_ALIAS  # type: ignore
+            )
+            .get_order_by()
+        ]
+        queryset = queryset.annotate(
+            _strawberry_row_number_reversed=_PaginationWindow(
+                RowNumber(),
+                partition_by=related_field_id,
+                order_by=order_by_reverse,
+            ),
         )
+        return queryset.filter(_strawberry_row_number_reversed__lte=limit)
 
     # Limit == -1 means no limit. sys.maxsize is set by relay when paginating
     # from the end to as a way to mimic a "not limit" as well

strawberry_django/relay/list_connection.py

@@ -9,11 +9,50 @@
 from strawberry.relay.types import NodeIterableType
 from strawberry.types import get_object_definition
 from strawberry.types.base import StrawberryContainer
+from strawberry.types.nodes import InlineFragment, Selection
 from strawberry.utils.await_maybe import AwaitableOrValue
+from strawberry.utils.inspect import in_async_context
 from typing_extensions import Self, deprecated
 
 from strawberry_django.pagination import get_total_count
+from strawberry_django.queryset import get_queryset_config
 from strawberry_django.resolvers import django_resolver
+from strawberry_django.utils.typing import unwrap_type
+
+
+def _should_optimize_total_count(info: Info) -> bool:
+    """Check if the user requested to resolve the `totalCount` field of a connection.
+
+    Taken and adjusted from strawberry.relay.utils
+    """
+    resolve_for_field_names = {"totalCount"}
+
+    def _check_selection(selection: Selection) -> bool:
+        """Recursively inspect the selection to check if the user requested to resolve the `edges` field.
+
+        Args:
+            selection (Selection): The selection to check.
+
+        Returns:
+            bool: True if the user requested to resolve the `edges` field of a connection, False otherwise.
+
+        """
+        if (
+            not isinstance(selection, InlineFragment)
+            and selection.name in resolve_for_field_names
+        ):
+            return True
+        if selection.selections:
+            return any(
+                _check_selection(selection) for selection in selection.selections
+            )
+        return False
+
+    for selection_field in info.selected_fields:
+        for selection in selection_field.selections:
+            if _check_selection(selection):
+                return True
+    return False
 
 
 @strawberry.type(name="Connection", description="A connection to a list of items.")
@@ -25,6 +64,11 @@ class DjangoListConnection(relay.ListConnection[relay.NodeType]):
     def total_count(self) -> Optional[int]:
         assert self.nodes is not None
 
+        try:
+            return self.edges[0].node._strawberry_total_count  # type: ignore
+        except (IndexError, AttributeError):
+            pass
+
         if isinstance(self.nodes, models.QuerySet):
             return get_total_count(self.nodes)
 
@@ -42,32 +86,52 @@ def resolve_connection(
         last: Optional[int] = None,
         **kwargs: Any,
     ) -> AwaitableOrValue[Self]:
-        from strawberry_django.optimizer import is_optimized_by_prefetching
+        if isinstance(nodes, models.QuerySet) and (
+            queryset_config := get_queryset_config(nodes)
+        ):
+            if queryset_config.optimized_by_prefetching:
+                try:
+                    conn = cls.resolve_optimized_connection_by_prefetch(
+                        nodes,
+                        info=info,
+                        before=before,
+                        after=after,
+                        first=first,
+                        last=last,
+                        **kwargs,
+                    )
+                except AttributeError:
+                    warnings.warn(
+                        (
+                            "Pagination annotations not found, falling back to QuerySet resolution. "
+                            "This might cause N+1 issues..."
+                        ),
+                        RuntimeWarning,
+                        stacklevel=2,
+                    )
+                else:
+                    conn = cast("Self", conn)
+                    conn.nodes = nodes
+                    return conn
 
-        if isinstance(nodes, models.QuerySet) and is_optimized_by_prefetching(nodes):
-            try:
-                conn = cls.resolve_connection_from_cache(
-                    nodes,
-                    info=info,
-                    before=before,
-                    after=after,
-                    first=first,
-                    last=last,
-                    **kwargs,
-                )
-            except AttributeError:
-                warnings.warn(
-                    (
-                        "Pagination annotations not found, falling back to QuerySet resolution. "
-                        "This might cause N+1 issues..."
-                    ),
-                    RuntimeWarning,
-                    stacklevel=2,
-                )
-            else:
-                conn = cast("Self", conn)
-                conn.nodes = nodes
-                return conn
+            if queryset_config.optimized:
+                if (last or 0) > 0 and before is None:
+                    return cls.resolve_optimized_last_connection(
+                        nodes,
+                        info=info,
+                        before=before,
+                        after=after,
+                        first=first,
+                        last=last,
+                        **kwargs,
+                    )
+
+                if _should_optimize_total_count(info):
+                    nodes = nodes.annotate(
+                        _strawberry_total_count=models.Window(
+                            expression=models.Count(1), partition_by=None
+                        )
+                    )
 
         conn = super().resolve_connection(
             nodes,
@@ -93,7 +157,7 @@ async def wrapper():
         return conn
 
     @classmethod
-    def resolve_connection_from_cache(
+    def resolve_optimized_connection_by_prefetch(
         cls,
         nodes: NodeIterableType[relay.NodeType],
         *,
@@ -147,6 +211,63 @@ def resolve_connection_from_cache(
             ),
         )
 
+    @classmethod
+    def resolve_optimized_last_connection(
+        cls,
+        nodes: NodeIterableType[relay.NodeType],
+        *,
+        info: Info,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        first: Optional[int] = None,
+        last: Optional[int] = None,
+        **kwargs: Any,
+    ) -> AwaitableOrValue[Self]:
+        """Resolve the connection being paginated only via `last`.
+
+        In order to prevent fetching the entire table, QuerySet is first counted & the
+        amount is used instead of `before=None`.
+        """
+        assert isinstance(nodes, models.QuerySet)
+
+        type_def = get_object_definition(cls)
+        assert type_def
+        field_def = type_def.get_field("edges")
+        assert field_def
+        field = field_def.resolve_type(type_definition=type_def)
+        field = unwrap_type(field)
+        edge_class = cast("relay.Edge[relay.NodeType]", field)
+
+        if in_async_context():
+
+            async def wrapper():
+                total_count = await nodes.acount()
+                before = relay.to_base64(edge_class.CURSOR_PREFIX, total_count)
+                conn = cls.resolve_connection(
+                    nodes,
+                    info=info,
+                    before=before,
+                    after=after,
+                    first=first,
+                    last=last,
+                    **kwargs,
+                )
+                return await conn if inspect.isawaitable(conn) else conn
+
+            return wrapper()
+
+        total_count = nodes.count()
+        before = relay.to_base64(edge_class.CURSOR_PREFIX, total_count)
+        return cls.resolve_connection(
+            nodes,
+            info=info,
+            before=before,
+            after=after,
+            first=first,
+            last=last,
+            **kwargs,
+        )
+
 
 if TYPE_CHECKING:
 

tests/polymorphism_custom/test_optimizer.py

@@ -186,7 +186,7 @@ def test_polymorphic_interface_connection():
     }
     """
 
-    with assert_num_queries(2):
+    with assert_num_queries(1):
         result = schema.execute_sync(query)
     assert not result.errors
     assert result.data == {