incremental: allow nested defers at the same level

Replicates graphql/graphql-js@51f41eb

Author: Cito

docs/conf.py

@@ -184,6 +184,7 @@
     "graphql.execution.build_field_plan.FieldPlan",
     "graphql.execution.collect_fields.DeferUsage",
     "graphql.execution.execute.StreamArguments",
+    "graphql.execution.execute.SubFieldPlan",
     "graphql.execution.execute.StreamUsage",
     "graphql.execution.map_async_iterable.map_async_iterable",
     "graphql.execution.incremental_publisher.CompletedResult",

src/graphql/execution/build_field_plan.py

@@ -34,7 +34,6 @@ class FieldGroup(NamedTuple):
 
     fields: list[FieldDetails]
     defer_usages: DeferUsageSet | None = None
-    known_defer_usages: DeferUsageSet | None = None
 
     def to_nodes(self) -> list[FieldNode]:
         """Return the field nodes in this group."""
@@ -59,22 +58,15 @@ class FieldPlan(NamedTuple):
 
     grouped_field_set: GroupedFieldSet
     new_grouped_field_set_details_map: RefMap[DeferUsageSet, NewGroupedFieldSetDetails]
-    new_defer_usages: list[DeferUsage]
 
 
 def build_field_plan(
     fields: dict[str, list[FieldDetails]],
     parent_defer_usages: DeferUsageSet | None = None,
-    known_defer_usages: DeferUsageSet | None = None,
 ) -> FieldPlan:
     """Build a plan for executing fields."""
     if parent_defer_usages is None:
         parent_defer_usages = RefSet()
-    if known_defer_usages is None:
-        known_defer_usages = RefSet()
-
-    new_defer_usages: RefSet[DeferUsage] = RefSet()
-    new_known_defer_usages: RefSet[DeferUsage] = RefSet(known_defer_usages)
 
     grouped_field_set: GroupedFieldSet = {}
 
@@ -93,9 +85,6 @@ def build_field_plan(
                 in_original_result = True
                 continue
             defer_usage_set.add(defer_usage)
-            if defer_usage not in known_defer_usages:
-                new_defer_usages.add(defer_usage)
-                new_known_defer_usages.add(defer_usage)
         if in_original_result:
             defer_usage_set.clear()
         else:
@@ -112,7 +101,7 @@ def build_field_plan(
         if defer_usage_set == parent_defer_usages:
             field_group = grouped_field_set.get(response_key)
             if field_group is None:  # pragma: no cover else
-                field_group = FieldGroup([], defer_usage_set, new_known_defer_usages)
+                field_group = FieldGroup([], defer_usage_set)
                 grouped_field_set[response_key] = field_group
             field_group.fields.extend(field_details_list)
             continue
@@ -139,10 +128,8 @@ def build_field_plan(
 
         field_group = new_grouped_field_set.get(response_key)
         if field_group is None:  # pragma: no cover else
-            field_group = FieldGroup([], defer_usage_set, new_known_defer_usages)
+            field_group = FieldGroup([], defer_usage_set)
             new_grouped_field_set[response_key] = field_group
         field_group.fields.extend(field_details_list)
 
-    return FieldPlan(
-        grouped_field_set, new_grouped_field_set_details_map, list(new_defer_usages)
-    )
+    return FieldPlan(grouped_field_set, new_grouped_field_set_details_map)

src/graphql/execution/collect_fields.py

@@ -27,6 +27,7 @@
 
 __all__ = [
     "CollectFieldsContext",
+    "CollectedFields",
     "DeferUsage",
     "FieldDetails",
     "collect_fields",
@@ -69,13 +70,20 @@ class CollectFieldsContext(NamedTuple):
     visited_fragment_names: set[str]
 
 
+class CollectedFields(NamedTuple):
+    """Collected fields with new defer usages."""
+
+    fields: dict[str, list[FieldDetails]]
+    new_defer_usages: list[DeferUsage]
+
+
 def collect_fields(
     schema: GraphQLSchema,
     fragments: dict[str, FragmentDefinitionNode],
     variable_values: dict[str, Any],
     runtime_type: GraphQLObjectType,
     operation: OperationDefinitionNode,
-) -> dict[str, list[FieldDetails]]:
+) -> CollectedFields:
     """Collect fields.
 
     Given a selection_set, collects all the fields and returns them.
@@ -87,6 +95,7 @@ def collect_fields(
     For internal use only.
     """
     grouped_field_set: dict[str, list[FieldDetails]] = defaultdict(list)
+    new_defer_usages: list[DeferUsage] = []
     context = CollectFieldsContext(
         schema,
         fragments,
@@ -96,8 +105,10 @@ def collect_fields(
         set(),
     )
 
-    collect_fields_impl(context, operation.selection_set, grouped_field_set)
-    return grouped_field_set
+    collect_fields_impl(
+        context, operation.selection_set, grouped_field_set, new_defer_usages
+    )
+    return CollectedFields(grouped_field_set, new_defer_usages)
 
 
 def collect_subfields(
@@ -107,7 +118,7 @@ def collect_subfields(
     operation: OperationDefinitionNode,
     return_type: GraphQLObjectType,
     field_details: list[FieldDetails],
-) -> dict[str, list[FieldDetails]]:
+) -> CollectedFields:
     """Collect subfields.
 
     Given a list of field nodes, collects all the subfields of the passed in fields,
@@ -128,6 +139,7 @@ def collect_subfields(
         set(),
     )
     sub_grouped_field_set: dict[str, list[FieldDetails]] = defaultdict(list)
+    new_defer_usages: list[DeferUsage] = []
 
     for field_detail in field_details:
         node = field_detail.node
@@ -136,17 +148,18 @@ def collect_subfields(
                 context,
                 node.selection_set,
                 sub_grouped_field_set,
+                new_defer_usages,
                 field_detail.defer_usage,
             )
 
-    return sub_grouped_field_set
+    return CollectedFields(sub_grouped_field_set, new_defer_usages)
 
 
 def collect_fields_impl(
     context: CollectFieldsContext,
     selection_set: SelectionSetNode,
     grouped_field_set: dict[str, list[FieldDetails]],
-    parent_defer_usage: DeferUsage | None = None,
+    new_defer_usages: list[DeferUsage],
     defer_usage: DeferUsage | None = None,
 ) -> None:
     """Collect fields (internal implementation)."""
@@ -164,31 +177,39 @@ def collect_fields_impl(
             if not should_include_node(variable_values, selection):
                 continue
             key = get_field_entry_key(selection)
-            grouped_field_set[key].append(
-                FieldDetails(selection, defer_usage or parent_defer_usage)
-            )
+            grouped_field_set[key].append(FieldDetails(selection, defer_usage))
         elif isinstance(selection, InlineFragmentNode):
             if not should_include_node(
                 variable_values, selection
             ) or not does_fragment_condition_match(schema, selection, runtime_type):
                 continue
 
             new_defer_usage = get_defer_usage(
-                operation, variable_values, selection, parent_defer_usage
+                operation, variable_values, selection, defer_usage
             )
 
-            collect_fields_impl(
-                context,
-                selection.selection_set,
-                grouped_field_set,
-                parent_defer_usage,
-                new_defer_usage or defer_usage,
-            )
+            if new_defer_usage is None:
+                collect_fields_impl(
+                    context,
+                    selection.selection_set,
+                    grouped_field_set,
+                    new_defer_usages,
+                    defer_usage,
+                )
+            else:
+                new_defer_usages.append(new_defer_usage)
+                collect_fields_impl(
+                    context,
+                    selection.selection_set,
+                    grouped_field_set,
+                    new_defer_usages,
+                    new_defer_usage,
+                )
         elif isinstance(selection, FragmentSpreadNode):  # pragma: no cover else
             frag_name = selection.name.value
 
             new_defer_usage = get_defer_usage(
-                operation, variable_values, selection, parent_defer_usage
+                operation, variable_values, selection, defer_usage
             )
 
             if new_defer_usage is None and (
@@ -205,14 +226,22 @@ def collect_fields_impl(
 
             if new_defer_usage is None:
                 visited_fragment_names.add(frag_name)
-
-            collect_fields_impl(
-                context,
-                fragment.selection_set,
-                grouped_field_set,
-                parent_defer_usage,
-                new_defer_usage or defer_usage,
-            )
+                collect_fields_impl(
+                    context,
+                    fragment.selection_set,
+                    grouped_field_set,
+                    new_defer_usages,
+                    defer_usage,
+                )
+            else:
+                new_defer_usages.append(new_defer_usage)
+                collect_fields_impl(
+                    context,
+                    fragment.selection_set,
+                    grouped_field_set,
+                    new_defer_usages,
+                    new_defer_usage,
+                )
 
 
 def get_defer_usage(

src/graphql/execution/execute.py

@@ -79,7 +79,6 @@
 from .build_field_plan import (
     DeferUsageSet,
     FieldGroup,
-    FieldPlan,
     GroupedFieldSet,
     NewGroupedFieldSetDetails,
     build_field_plan,
@@ -156,6 +155,14 @@ class StreamUsage(NamedTuple):
     field_group: FieldGroup
 
 
+class SubFieldPlan(NamedTuple):
+    """A plan for executing fields with defer usages."""
+
+    grouped_field_set: GroupedFieldSet
+    new_grouped_field_set_details_map: RefMap[DeferUsageSet, NewGroupedFieldSetDetails]
+    new_defer_usages: list[DeferUsage]
+
+
 class ExecutionContext:
     """Data that must be available at all points during query execution.
 
@@ -206,7 +213,7 @@ def __init__(
         if is_awaitable:
             self.is_awaitable = is_awaitable
         self._canceled_iterators: set[AsyncIterator] = set()
-        self._field_plan_cache: dict[tuple, FieldPlan] = {}
+        self._sub_field_plan_cache: dict[tuple, SubFieldPlan] = {}
         self._tasks: set[Awaitable] = set()
         self._stream_usages: RefMap[FieldGroup, StreamUsage] = RefMap()
 
@@ -324,12 +331,10 @@ def execute_operation(
             )
             raise GraphQLError(msg, operation)
 
-        fields = collect_fields(
+        fields, new_defer_usages = collect_fields(
             schema, self.fragments, self.variable_values, root_type, operation
         )
-        grouped_field_set, new_grouped_field_set_details_map, new_defer_usages = (
-            build_field_plan(fields)
-        )
+        grouped_field_set, new_grouped_field_set_details_map = build_field_plan(fields)
 
         incremental_publisher = self.incremental_publisher
         new_defer_map = add_new_deferred_fragments(
@@ -1313,14 +1318,14 @@ def collect_and_execute_subfields(
 
     def build_sub_field_plan(
         self, return_type: GraphQLObjectType, field_group: FieldGroup
-    ) -> FieldPlan:
+    ) -> SubFieldPlan:
         """Collect subfields.
 
         A memoized function for building subfield plans with regard to the return type.
         Memoizing ensures the  subfields are not repeatedly calculated, which saves
         overhead when resolving lists of values.
         """
-        cache = self._field_plan_cache
+        cache = self._sub_field_plan_cache
         # We cannot use the field_group itself as key for the cache, since it
         # is not hashable as a list. We also do not want to use the field_group
         # itself (converted to a tuple) as keys, since hashing them is slow.
@@ -1333,21 +1338,20 @@ def build_sub_field_plan(
             if len(field_group) == 1  # optimize most frequent case
             else (return_type, *map(id, field_group))
         )
-        plan = cache.get(key)
-        if plan is None:
-            sub_fields = collect_subfields(
+        sub_field_plan = cache.get(key)
+        if sub_field_plan is None:
+            sub_fields, new_defer_usages = collect_subfields(
                 self.schema,
                 self.fragments,
                 self.variable_values,
                 self.operation,
                 return_type,
                 field_group.fields,
             )
-            plan = build_field_plan(
-                sub_fields, field_group.defer_usages, field_group.known_defer_usages
-            )
-            cache[key] = plan
-        return plan
+            field_plan = build_field_plan(sub_fields, field_group.defer_usages)
+            sub_field_plan = SubFieldPlan(*field_plan, new_defer_usages)
+            cache[key] = sub_field_plan
+        return sub_field_plan
 
     def map_source_to_response(
         self, result_or_stream: ExecutionResult | AsyncIterable[Any]
@@ -2301,7 +2305,7 @@ def execute_subscription(
         context.variable_values,
         root_type,
         context.operation,
-    )
+    ).fields
 
     first_root_field = next(iter(fields.items()))
     response_name, field_details_list = first_root_field

src/graphql/validation/rules/single_field_subscriptions.py

@@ -52,7 +52,7 @@ def enter_operation_definition(
                 variable_values,
                 subscription_type,
                 node,
-            )
+            ).fields
             if len(fields) > 1:
                 field_groups = list(fields.values())
                 extra_field_groups = field_groups[1:]