Skip to content

Commit fd9edfb

Browse files
committed
Await values in execute_stream_async_iterator
1 parent 9627bc5 commit fd9edfb

File tree

2 files changed

+40
-4
lines changed

2 files changed

+40
-4
lines changed

src/graphql/execution/execute.py

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1682,6 +1682,7 @@ async def execute_stream_async_iterator(
16821682
parent_context: IncrementalDataRecord | None = None,
16831683
) -> None:
16841684
"""Execute stream iterator."""
1685+
is_awaitable = self.is_awaitable
16851686
incremental_publisher = self.incremental_publisher
16861687
index = initial_index
16871688
previous_incremental_data_record = parent_context
@@ -1721,10 +1722,35 @@ async def execute_stream_async_iterator(
17211722
except StopAsyncIteration:
17221723
done = True
17231724

1724-
incremental_publisher.complete_stream_items_record(
1725-
incremental_data_record,
1726-
[completed_item],
1727-
)
1725+
if is_awaitable(completed_item):
1726+
1727+
async def await_completed_item(
1728+
incremental_data_record: StreamItemsRecord, completed_item: Any
1729+
) -> None:
1730+
try:
1731+
value = await completed_item
1732+
except GraphQLError as error:
1733+
incremental_publisher.add_field_error(
1734+
incremental_data_record, error
1735+
)
1736+
incremental_publisher.filter(path, incremental_data_record)
1737+
incremental_publisher.complete_stream_items_record(
1738+
incremental_data_record, None
1739+
)
1740+
else:
1741+
incremental_publisher.complete_stream_items_record(
1742+
incremental_data_record, [value]
1743+
)
1744+
1745+
self.add_task(
1746+
await_completed_item(incremental_data_record, completed_item)
1747+
)
1748+
1749+
else:
1750+
incremental_publisher.complete_stream_items_record(
1751+
incremental_data_record,
1752+
[completed_item],
1753+
)
17281754

17291755
if done:
17301756
break

tests/execution/test_stream.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1708,6 +1708,9 @@ async def friend_list(_info):
17081708
],
17091709
},
17101710
],
1711+
"hasNext": True,
1712+
},
1713+
{
17111714
"hasNext": False,
17121715
},
17131716
]
@@ -1848,6 +1851,9 @@ async def get_friends(_info):
18481851
"path": ["friendList", 2],
18491852
}
18501853
],
1854+
"hasNext": True,
1855+
},
1856+
{
18511857
"hasNext": False,
18521858
},
18531859
]
@@ -1927,6 +1933,10 @@ async def get_friends(_info):
19271933
"path": ["nestedObject", "nestedFriendList", 1],
19281934
},
19291935
],
1936+
"hasNext": True,
1937+
}
1938+
result5 = await anext(iterator)
1939+
assert result5.formatted == {
19301940
"hasNext": False,
19311941
}
19321942

0 commit comments

Comments
 (0)