refactor(incremental): enqueue only released records

Replicates graphql/graphql-js@0732a87

Author: Cito

src/graphql/execution/incremental_graph.py

@@ -26,12 +26,10 @@
     from graphql.execution.types import (
         DeferredFragmentRecord,
         DeferredGroupedFieldSetRecord,
-        DeferredGroupedFieldSetResult,
         IncrementalDataRecord,
         IncrementalDataRecordResult,
         ReconcilableDeferredGroupedFieldSetResult,
         StreamItemsRecord,
-        StreamItemsResult,
         SubsequentResultRecord,
     )
 
@@ -51,20 +49,17 @@ class DeferredFragmentNode:
         "deferred_fragment_record",
         "deferred_grouped_field_set_records",
         "reconcilable_results",
-        "results",
     )
 
     deferred_fragment_record: DeferredFragmentRecord
     deferred_grouped_field_set_records: dict[DeferredGroupedFieldSetRecord, None]
-    results: list[DeferredGroupedFieldSetResult]
     reconcilable_results: dict[ReconcilableDeferredGroupedFieldSetResult, None]
     children: list[DeferredFragmentNode]
 
     def __init__(self, deferred_fragment_record: DeferredFragmentRecord) -> None:
         """Initialize the DeferredFragmentNode."""
         self.deferred_fragment_record = deferred_fragment_record
         self.deferred_grouped_field_set_records = {}
-        self.results = []
         self.reconcilable_results = {}
         self.children = []
 
@@ -95,6 +90,7 @@ class IncrementalGraph:
     _pending: dict[SubsequentResultNode, None]
     _deferred_fragment_nodes: dict[DeferredFragmentRecord, DeferredFragmentNode]
     _new_pending: dict[SubsequentResultNode, None]
+    _new_incremental_data_records: dict[IncrementalDataRecord, None]
     _completed_queue: list[IncrementalDataRecordResult]
     _next_queue: list[Future[Iterable[IncrementalDataRecordResult]]]
 
@@ -105,6 +101,7 @@ def __init__(self) -> None:
         self._pending = {}
         self._deferred_fragment_nodes = {}
         self._new_pending = {}
+        self._new_incremental_data_records = {}
         self._completed_queue = []
         self._next_queue = []
         self._tasks = set()
@@ -115,49 +112,10 @@ def add_incremental_data_records(
         """Add incremental data records."""
         for incremental_data_record in incremental_data_records:
             if is_deferred_grouped_field_set_record(incremental_data_record):
-                for deferred_fragment_record in (
-                    incremental_data_record.deferred_fragment_records
-                ):  # pragma: no branch
-                    deferred_fragment_node = self._add_deferred_fragment_node(
-                        deferred_fragment_record
-                    )
-                    deferred_fragment_node.deferred_grouped_field_set_records[
-                        incremental_data_record
-                    ] = None
-
-                deferred_result = incremental_data_record.result
-                if is_awaitable(deferred_result):
-
-                    async def enqueue_deferred(
-                        deferred_result: Awaitable[DeferredGroupedFieldSetResult],
-                    ) -> None:
-                        self._enqueue_completed_deferred_grouped_field_set(
-                            await deferred_result
-                        )
-
-                    self._add_task(enqueue_deferred(deferred_result))
-                else:
-                    self._enqueue_completed_deferred_grouped_field_set(
-                        deferred_result,  # type: ignore
-                    )
-                continue
-
-            incremental_data_record = cast("StreamItemsRecord", incremental_data_record)
-            stream_record = incremental_data_record.stream_record
-            if stream_record.id is None:
-                self._new_pending[stream_record] = None
-
-            stream_result = incremental_data_record.result
-            if is_awaitable(stream_result):
-
-                async def enqueue_stream(
-                    stream_result: Awaitable[StreamItemsResult],
-                ) -> None:
-                    self._enqueue(await stream_result)
-
-                self._add_task(enqueue_stream(stream_result))
+                self._add_deferred_grouped_field_set_record(incremental_data_record)
             else:
-                self._enqueue(stream_result)  # type: ignore
+                stream_items_record = cast("StreamItemsRecord", incremental_data_record)
+                self._add_stream_items_record(stream_items_record)
 
     def add_completed_reconcilable_deferred_grouped_field_set(
         self, reconcilable_result: ReconcilableDeferredGroupedFieldSetResult
@@ -178,6 +136,7 @@ def add_completed_reconcilable_deferred_grouped_field_set(
     def get_new_pending(self) -> list[SubsequentResultRecord]:
         """Get new pending subsequent result records."""
         _pending, _new_pending = self._pending, self._new_pending
+        _new_incremental_data_records = self._new_incremental_data_records
         new_pending: list[SubsequentResultRecord] = []
         add_result = new_pending.append
         # avoid iterating over a changing dict
@@ -189,13 +148,34 @@ def get_new_pending(self) -> list[SubsequentResultRecord]:
                 _pending[node] = None
                 add_result(node)
             elif node.deferred_grouped_field_set_records:  # type: ignore
+                records = node.deferred_grouped_field_set_records  # type: ignore
+                for deferred_grouped_field_set_node in records:  # pragma: no branch
+                    _new_incremental_data_records[deferred_grouped_field_set_node] = (
+                        None
+                    )
                 _pending[node] = None
                 add_result(node.deferred_fragment_record)  # type: ignore
             else:
                 for child in node.children:  # type: ignore
                     _new_pending[child] = None
                     add_iteration(child)
         _new_pending.clear()
+
+        enqueue = self._enqueue
+        for incremental_data_record in _new_incremental_data_records:
+            result = incremental_data_record.result
+            if is_awaitable(result):
+
+                async def enqueue_incremental(
+                    result: Awaitable[IncrementalDataRecordResult],
+                ) -> None:
+                    enqueue(await result)
+
+                self._add_task(enqueue_incremental(result))
+            else:
+                enqueue(result)  # type: ignore
+        _new_incremental_data_records.clear()
+
         return new_pending
 
     async def completed_incremental_data(
@@ -247,8 +227,6 @@ def complete_deferred_fragment(
         new_pending = self._new_pending
         for child in deferred_fragment_node.children:
             new_pending[child] = None
-            for result in child.results:
-                self._enqueue(result)
         return reconcilable_results
 
     def remove_deferred_fragment(
@@ -277,6 +255,29 @@ def _remove_pending(self, subsequent_result_node: SubsequentResultNode) -> None:
         if not self._pending:
             self.stop_incremental_data()
 
+    def _add_deferred_grouped_field_set_record(
+        self, deferred_grouped_field_set_record: DeferredGroupedFieldSetRecord
+    ) -> None:
+        """Add a deferred grouped field set record."""
+        pending = self._pending
+        new_incremental_data_records = self._new_incremental_data_records
+        add = self._add_deferred_fragment_node
+        records = deferred_grouped_field_set_record.deferred_fragment_records
+        for deferred_fragment_record in records:  # pragma: no branch
+            deferred_fragment_node = add(deferred_fragment_record)
+            if deferred_fragment_node in pending:
+                new_incremental_data_records[deferred_grouped_field_set_record] = None
+            deferred_fragment_node.deferred_grouped_field_set_records[
+                deferred_grouped_field_set_record
+            ] = None
+
+    def _add_stream_items_record(self, stream_items_record: StreamItemsRecord) -> None:
+        """Add a stream items record."""
+        stream_record = stream_items_record.stream_record
+        if stream_record not in self._pending:
+            self._new_pending[stream_record] = None
+        self._new_incremental_data_records[stream_items_record] = None
+
     def _add_deferred_fragment_node(
         self, deferred_fragment_record: DeferredFragmentRecord
     ) -> DeferredFragmentNode:
@@ -298,24 +299,6 @@ def _add_deferred_fragment_node(
                 parent_node.children.append(deferred_fragment_node)
         return deferred_fragment_node
 
-    def _enqueue_completed_deferred_grouped_field_set(
-        self, result: DeferredGroupedFieldSetResult
-    ) -> None:
-        """Enqueue completed deferred grouped field set result."""
-        is_pending = False
-        nodes = self._deferred_fragment_nodes
-        record = result.deferred_grouped_field_set_record
-        for deferred_fragment_record in record.deferred_fragment_records:
-            try:
-                deferred_fragment_node = nodes[deferred_fragment_record]
-            except KeyError:  # pragma: no cover
-                continue
-            if deferred_fragment_node in self._pending:
-                is_pending = True
-            deferred_fragment_node.results.append(result)
-        if is_pending:
-            self._enqueue(result)
-
     def _add_task(self, awaitable: Awaitable[Any]) -> None:
         """Add the given task to the tasks set for later execution."""
         tasks = self._tasks

tests/execution/test_defer.py

@@ -726,6 +726,12 @@ async def can_defer_a_fragment_within_an_already_deferred_fragment():
                 "pending": [{"id": "1", "path": ["hero"], "label": "DeferNested"}],
                 "incremental": [
                     {"data": {"id": "1"}, "id": "0"},
+                ],
+                "completed": [{"id": "0"}],
+                "hasNext": True,
+            },
+            {
+                "incremental": [
                     {
                         "data": {
                             "friends": [
@@ -737,7 +743,7 @@ async def can_defer_a_fragment_within_an_already_deferred_fragment():
                         "id": "1",
                     },
                 ],
-                "completed": [{"id": "0"}, {"id": "1"}],
+                "completed": [{"id": "1"}],
                 "hasNext": False,
             },
         ]
@@ -1324,22 +1330,39 @@ async def deduplicates_fields_with_deferred_fragments_at_multiple_levels():
             {
                 "pending": [
                     {"id": "1", "path": ["hero", "nestedObject"]},
-                    {"id": "2", "path": ["hero", "nestedObject", "deeperObject"]},
                 ],
                 "incremental": [
                     {
                         "data": {"bar": "bar"},
                         "id": "0",
                         "subPath": ["nestedObject", "deeperObject"],
                     },
+                ],
+                "completed": [{"id": "0"}],
+                "hasNext": True,
+            },
+            {
+                "pending": [
+                    {"id": "2", "path": ["hero", "nestedObject", "deeperObject"]}
+                ],
+                "incremental": [
                     {
                         "data": {"baz": "baz"},
                         "id": "1",
                         "subPath": ["deeperObject"],
                     },
-                    {"data": {"bak": "bak"}, "id": "2"},
                 ],
-                "completed": [{"id": "0"}, {"id": "1"}, {"id": "2"}],
+                "completed": [{"id": "1"}],
+                "hasNext": True,
+            },
+            {
+                "incremental": [
+                    {
+                        "data": {"bak": "bak"},
+                        "id": "2",
+                    },
+                ],
+                "completed": [{"id": "2"}],
                 "hasNext": False,
             },
         ]

tests/execution/test_stream.py

@@ -1624,6 +1624,10 @@ async def friend_list(_info):
             },
             {
                 "incremental": [
+                    {
+                        "items": [{"name": "Luke"}],
+                        "id": "1",
+                    },
                     {
                         "data": {"scalarField": None},
                         "id": "0",
@@ -1635,7 +1639,6 @@ async def friend_list(_info):
                             },
                         ],
                     },
-                    {"items": [{"name": "Luke"}], "id": "1"},
                 ],
                 "completed": [{"id": "0"}],
                 "hasNext": True,
@@ -2082,8 +2085,8 @@ async def get_friends(_info):
         assert result2.formatted == {
             "pending": [{"id": "2", "path": ["friendList", 1], "label": "DeferName"}],
             "incremental": [
-                {"data": {"name": "Luke"}, "id": "0"},
                 {"items": [{"id": "2"}], "id": "1"},
+                {"data": {"name": "Luke"}, "id": "0"},
             ],
             "completed": [{"id": "0"}],
             "hasNext": True,
@@ -2160,8 +2163,8 @@ async def get_friends(_info):
         assert result2.formatted == {
             "pending": [{"id": "2", "path": ["friendList", 1], "label": "DeferName"}],
             "incremental": [
-                {"data": {"name": "Luke"}, "id": "0"},
                 {"items": [{"id": "2"}], "id": "1"},
+                {"data": {"name": "Luke"}, "id": "0"},
             ],
             "completed": [{"id": "0"}],
             "hasNext": True,
@@ -2233,8 +2236,8 @@ async def iterable(_info):
         assert result2 == {
             "pending": [{"id": "2", "path": ["friendList", 1]}],
             "incremental": [
-                {"data": {"name": "Luke"}, "id": "0"},
                 {"items": [{"id": "2"}], "id": "1"},
+                {"data": {"name": "Luke"}, "id": "0"},
             ],
             "completed": [{"id": "0"}],
             "hasNext": True,