diff --git a/docs/modules/pyutils.rst b/docs/modules/pyutils.rst index e8ce0407..92ac98a4 100644 --- a/docs/modules/pyutils.rst +++ b/docs/modules/pyutils.rst @@ -24,7 +24,6 @@ PyUtils .. autoclass:: AbortSignal .. autoexception:: AbortError .. autoclass:: AwaitableOrValue -.. autoclass:: BoxedAwaitableOrValue .. autofunction:: suggestion_list .. autoclass:: FrozenError :no-members: diff --git a/pyproject.toml b/pyproject.toml index 585c6af7..efa5319b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -149,6 +149,7 @@ ignore = [ "I001", # imports do not need to be sorted ] "src/graphql/execution/*" = [ + "ARG002", # allow unused arguments in overridable hook methods "BLE001", # allow catching blind exception ] "src/graphql/language/ast.py" = [ diff --git a/src/graphql/execution/__init__.py b/src/graphql/execution/__init__.py index dc9b7d8e..fff930a8 100644 --- a/src/graphql/execution/__init__.py +++ b/src/graphql/execution/__init__.py @@ -8,7 +8,6 @@ from .async_iterables import map_async_iterable from .types import ( CompletedResult, - DeferredFragmentRecord, ExecutionResult, ExperimentalIncrementalExecutionResults, FormattedSubsequentIncrementalExecutionResult, @@ -24,10 +23,6 @@ InitialIncrementalExecutionResult, PendingResult, SubsequentIncrementalExecutionResult, - StreamRecord, - StreamItemRecord, - StreamItemResult, - StreamItemsResult, ) from .middleware import MiddlewareManager from .values import ( @@ -50,7 +45,6 @@ AsyncWorkFinishedInfo, ExecutionHooks, Executor, - GraphQLWrappedResult, Middleware, RootSelectionSetExecutor, ) @@ -59,7 +53,6 @@ "AbortedGraphQLExecutionError", "AsyncWorkFinishedInfo", "CompletedResult", - "DeferredFragmentRecord", "ExecutionHooks", "ExecutionResult", "Executor", @@ -71,7 +64,6 @@ "FormattedInitialIncrementalExecutionResult", "FormattedPendingResult", "FormattedSubsequentIncrementalExecutionResult", - "GraphQLWrappedResult", "IncrementalDeferResult", "IncrementalResult", "IncrementalStreamResult", @@ -80,10 +72,6 @@ "MiddlewareManager", "PendingResult", "RootSelectionSetExecutor", - "StreamItemRecord", - "StreamItemResult", - "StreamItemsResult", - "StreamRecord", "SubsequentIncrementalExecutionResult", "VariableValues", "create_source_event_stream", diff --git a/src/graphql/execution/collect_fields.py b/src/graphql/execution/collect_fields.py index 59ab12e7..01e817a8 100644 --- a/src/graphql/execution/collect_fields.py +++ b/src/graphql/execution/collect_fields.py @@ -12,7 +12,6 @@ FragmentSpreadNode, InlineFragmentNode, OperationDefinitionNode, - OperationType, SelectionSetNode, ) from ..type import ( @@ -201,7 +200,7 @@ def collect_fields_impl( schema, fragments, variable_values, - operation, + _operation, runtime_type, visited_fragment_names, hide_suggestions, @@ -226,7 +225,6 @@ def collect_fields_impl( continue new_defer_usage = get_defer_usage( - operation, variable_values, fragment_variable_values, selection, @@ -267,7 +265,6 @@ def collect_fields_impl( continue new_defer_usage = get_defer_usage( - operation, variable_values, fragment_variable_values, selection, @@ -315,7 +312,6 @@ def collect_fields_impl( def get_defer_usage( - operation: OperationDefinitionNode, variable_values: VariableValues, fragment_variable_values: FragmentVariableValues | None, node: FragmentSpreadNode | InlineFragmentNode, @@ -334,13 +330,6 @@ def get_defer_usage( if not defer or defer.get("if") is False: return None - if operation.operation == OperationType.SUBSCRIPTION: - msg = ( - "`@defer` directive not supported on subscription operations." - " Disable `@defer` by setting the `if` argument to `false`." - ) - raise TypeError(msg) - return DeferUsage(defer.get("label"), parent_defer_usage) diff --git a/src/graphql/execution/execute.py b/src/graphql/execution/execute.py index 2d60e163..2f1c7fbe 100644 --- a/src/graphql/execution/execute.py +++ b/src/graphql/execution/execute.py @@ -2,59 +2,19 @@ from __future__ import annotations -from asyncio import ( - FIRST_COMPLETED, - Future, - TimeoutError, # only needed for Python < 3.11 # noqa: A004 - ensure_future, - gather, - get_running_loop, - iscoroutine, - sleep, - wait, -) -from collections.abc import ( - AsyncGenerator, - AsyncIterable, - AsyncIterator, - Awaitable, - Callable, - Iterable, - Iterator, - Mapping, - Sequence, -) -from contextlib import suppress -from copy import copy +from asyncio import ensure_future +from collections.abc import Callable from typing import ( TYPE_CHECKING, Any, - Generic, - NamedTuple, - TypeVar, cast, ) from ..error import GraphQLError, located_error -from ..language import ( - DocumentNode, - FieldNode, - FragmentDefinitionNode, - OperationDefinitionNode, - OperationType, - is_subscription_operation_definition_node, -) +from ..language import is_subscription_operation_definition_node from ..pyutils import ( - AbortSignal, - AwaitableOrValue, - BoxedAwaitableOrValue, Path, - RefMap, - Undefined, - async_reduce, - gather_with_cancel, inspect, - is_iterable, ) from ..pyutils.is_awaitable import ( is_async_iterable as default_is_async_iterable, @@ -62,83 +22,39 @@ from ..pyutils.is_awaitable import ( is_awaitable as default_is_awaitable, ) -from ..type import ( - GraphQLAbstractType, - GraphQLField, - GraphQLFieldResolver, - GraphQLLeafType, - GraphQLList, - GraphQLObjectType, - GraphQLOutputType, - GraphQLResolveInfo, - GraphQLResolveInfoHelpers, - GraphQLSchema, - GraphQLStreamDirective, - GraphQLTypeResolver, - assert_valid_schema, - is_abstract_type, - is_leaf_type, - is_list_type, - is_non_null_type, - is_object_type, -) -from ..type.directives import GraphQLDisableErrorPropagationDirective -from .aborted_graphql_execution_error import AbortedGraphQLExecutionError from .async_iterables import map_async_iterable -from .build_execution_plan import ( - DeferUsageSet, - ExecutionPlan, - FieldDetailsList, - GroupedFieldSet, - build_execution_plan, -) -from .collect_fields import ( - CollectedFields, - DeferUsage, - FieldDetails, - FragmentDetails, - collect_fields, - collect_subfields, -) -from .get_variable_signature import get_variable_signature -from .incremental_publisher import ( - IncrementalPublisherContext, - build_incremental_response, -) -from .middleware import MiddlewareManager -from .types import ( - CancellableStreamRecord, - CompletedExecutionGroup, - DeferredFragmentRecord, - ExecutionGroupResult, - ExecutionResult, - ExperimentalIncrementalExecutionResults, - FailedExecutionGroup, - IncrementalDataRecord, - PendingExecutionGroup, - StreamItemRecord, - StreamItemResult, - StreamRecord, - SuccessfulExecutionGroup, -) -from .values import ( - VariableValues, - get_argument_values, - get_directive_values, - get_variable_values, +from .collect_fields import collect_fields +from .executor import ( + UNEXPECTED_MULTIPLE_PAYLOADS, + AsyncWorkFinishedInfo, + ExecutionHooks, + Executor, + Middleware, + default_field_resolver, + default_type_resolver, ) +from .executor_throwing_on_incremental import ExecutorThrowingOnIncremental +from .incremental.incremental_executor import IncrementalExecutor +from .types import ExecutionResult, ExperimentalIncrementalExecutionResults +from .values import get_argument_values if TYPE_CHECKING: + from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable from typing import TypeAlias, TypeGuard - from ..pyutils import UndefinedType - from .get_variable_signature import GraphQLVariableSignature + from ..language import DocumentNode + from ..pyutils import AbortSignal, AwaitableOrValue + from ..type import ( + GraphQLFieldResolver, + GraphQLSchema, + GraphQLTypeResolver, + ) + from .middleware import MiddlewareManager __all__ = [ "AsyncWorkFinishedInfo", "ExecutionHooks", "Executor", - "GraphQLWrappedResult", "Middleware", "RootSelectionSetExecutor", "create_source_event_stream", @@ -153,9 +69,6 @@ "subscribe", ] -suppress_exceptions = suppress(Exception) -suppress_timeout_error = suppress(TimeoutError) - # Terminology # @@ -176,2284 +89,6 @@ # 3) inline fragment "spreads" e.g. "...on Type { a }" -Middleware: TypeAlias = tuple | list | MiddlewareManager | None - - -class StreamUsage(NamedTuple): - """Stream directive usage information""" - - label: str | None - initial_count: int - field_details_list: FieldDetailsList - - -class IncrementalContext: - """A context for incremental execution.""" - - errors: list[GraphQLError] | None - defer_usage_set: DeferUsageSet | None - - __slots__ = "defer_usage_set", "errors" - - def __init__(self, defer_usage_set: DeferUsageSet | None = None) -> None: - self.errors = None - self.defer_usage_set = defer_usage_set - - -class AsyncWorkFinishedInfo(NamedTuple): - """Information passed to the ``async_work_finished`` execution hook.""" - - executor: Executor - - -class ExecutionHooks(NamedTuple): - """Hooks for observing the execution of a GraphQL operation. - - The ``async_work_finished`` hook is run when all asynchronous work tracked - by the execution has finished. Cancelled asynchronous work may still be - running even after the result has been delivered; this hook allows - interested execution harnesses to track when this asynchronous work - completes. Errors raised by the hook are ignored. - """ - - async_work_finished: Callable[[AsyncWorkFinishedInfo], None] | None = None - - -class Executor(IncrementalPublisherContext): - """Data that must be available at all points during query execution. - - Namely, schema of the type system that is currently executing, and the fragments - defined in the query document. - """ - - schema: GraphQLSchema - # TODO: consider deprecating/removing fragment_definitions if/when fragment - # arguments are officially supported and/or the full fragment details are - # exposed within GraphQLResolveInfo. - fragment_definitions: dict[str, FragmentDefinitionNode] - fragments: dict[str, FragmentDetails] - root_value: Any - context_value: Any - operation: OperationDefinitionNode - variable_values: VariableValues - field_resolver: GraphQLFieldResolver - type_resolver: GraphQLTypeResolver - subscribe_field_resolver: GraphQLFieldResolver - enable_early_execution: bool - hide_suggestions: bool - abort_signal: AbortSignal | None - hooks: ExecutionHooks | None - async_helpers: GraphQLResolveInfoHelpers - errors: list[GraphQLError] | None - cancellable_streams: set[CancellableStreamRecord] | None - pending_incremental_futures: set[Future[Any]] - background_futures: set[Future[Any]] - async_work_finished_hook_task: Future[None] | None - middleware_manager: MiddlewareManager | None - error_propagation: bool - - is_awaitable: Callable[[Any], TypeGuard[Awaitable]] = staticmethod( - default_is_awaitable # type: ignore - ) - is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] = staticmethod( - default_is_async_iterable # type: ignore - ) - - def __init__( # noqa: PLR0913 - self, - schema: GraphQLSchema, - fragment_definitions: dict[str, FragmentDefinitionNode], - fragments: dict[str, FragmentDetails], - root_value: Any, - context_value: Any, - operation: OperationDefinitionNode, - variable_values: VariableValues, - field_resolver: GraphQLFieldResolver, - type_resolver: GraphQLTypeResolver, - subscribe_field_resolver: GraphQLFieldResolver, - enable_early_execution: bool = False, - middleware_manager: MiddlewareManager | None = None, - is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None, - is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None, - hide_suggestions: bool = False, - abort_signal: AbortSignal | None = None, - hooks: ExecutionHooks | None = None, - ) -> None: - self.schema = schema - self.fragment_definitions = fragment_definitions - self.fragments = fragments - self.root_value = root_value - self.context_value = context_value - self.operation = operation - self.variable_values = variable_values - self.field_resolver = field_resolver - self.type_resolver = type_resolver - self.subscribe_field_resolver = subscribe_field_resolver - self.enable_early_execution = enable_early_execution - self.hide_suggestions = hide_suggestions - self.abort_signal = abort_signal - self.hooks = hooks - self.async_helpers = GraphQLResolveInfoHelpers( - gather=self.gather_async_work, track=self.track_async_work - ) - self.middleware_manager = middleware_manager - self.error_propagation = not any( - directive.name.value == GraphQLDisableErrorPropagationDirective.name - for directive in operation.directives or () - ) - self.is_awaitable = is_awaitable or default_is_awaitable - self.is_async_iterable = is_async_iterable or default_is_async_iterable - self.errors = None - self.cancellable_streams = None - self.pending_incremental_futures = set() - self.background_futures = set() - self.async_work_finished_hook_task = None - self._relevant_sub_fields: dict[tuple, CollectedFields] = {} - self._stream_usages: RefMap[FieldDetailsList, StreamUsage] = RefMap() - self._execution_plans: RefMap[GroupedFieldSet, ExecutionPlan] = RefMap() - - @classmethod - def build( # noqa: PLR0913 - cls, - schema: GraphQLSchema, - document: DocumentNode, - root_value: Any = None, - context_value: Any = None, - raw_variable_values: dict[str, Any] | None = None, - operation_name: str | None = None, - field_resolver: GraphQLFieldResolver | None = None, - type_resolver: GraphQLTypeResolver | None = None, - subscribe_field_resolver: GraphQLFieldResolver | None = None, - max_coercion_errors: int = 50, - enable_early_execution: bool = False, - middleware: Middleware | None = None, - is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None, - is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None, - hide_suggestions: bool = False, - abort_signal: AbortSignal | None = None, - hooks: ExecutionHooks | None = None, - **custom_args: Any, - ) -> list[GraphQLError] | Executor: - """Build an executor - - Constructs an Executor object from the arguments passed to execute, which - we will pass throughout the other execution methods. - - Throws a GraphQLError if a valid executor cannot be created. - - For internal use only. - """ - # If the schema used for execution is invalid, raise an error. - assert_valid_schema(schema) - - operation: OperationDefinitionNode | None = None - fragment_definitions: dict[str, FragmentDefinitionNode] = {} - fragments: dict[str, FragmentDetails] = {} - fragment_variable_signature_errors: list[GraphQLError] = [] - middleware_manager: MiddlewareManager | None = None - if middleware is not None: - if isinstance(middleware, (list, tuple)): - middleware_manager = MiddlewareManager(*middleware) - elif isinstance(middleware, MiddlewareManager): - middleware_manager = middleware - else: - msg = ( - "Middleware must be passed as a list or tuple of functions" - " or objects, or as a single MiddlewareManager object." - f" Got {inspect(middleware)} instead." - ) - raise TypeError(msg) - - for definition in document.definitions: - if isinstance(definition, OperationDefinitionNode): - if operation_name is None: - if operation: - return [ - GraphQLError( - "Must provide operation name" - " if query contains multiple operations." - ) - ] - operation = definition - elif definition.name and definition.name.value == operation_name: - operation = definition - elif isinstance(definition, FragmentDefinitionNode): - fragment_definitions[definition.name.value] = definition - variable_signatures: dict[str, GraphQLVariableSignature] | None = None - if definition.variable_definitions: - variable_signatures = {} - for var_def in definition.variable_definitions: - signature = get_variable_signature(schema, var_def) - if isinstance(signature, GraphQLError): - fragment_variable_signature_errors.append(signature) - continue - variable_signatures[signature.name] = signature - fragments[definition.name.value] = FragmentDetails( - definition, variable_signatures - ) - - if not operation: - if operation_name is not None: - return [GraphQLError(f"Unknown operation named '{operation_name}'.")] - return [GraphQLError("Must provide an operation.")] - - if fragment_variable_signature_errors: - return fragment_variable_signature_errors - - variable_values = get_variable_values( - schema, - operation.variable_definitions or (), - raw_variable_values or {}, - max_errors=max_coercion_errors, - hide_suggestions=hide_suggestions, - ) - - if isinstance(variable_values, list): - return variable_values # errors - - return cls( - schema, - fragment_definitions, - fragments, - root_value, - context_value, - operation, - variable_values, - field_resolver or default_field_resolver, - type_resolver or default_type_resolver, - subscribe_field_resolver or default_field_resolver, - enable_early_execution, - middleware_manager, - is_awaitable, - is_async_iterable, - hide_suggestions=hide_suggestions, - abort_signal=abort_signal, - hooks=hooks, - **custom_args, - ) - - def build_per_event_executor(self, payload: Any) -> Executor: - """Create a copy of the executor for usage with subscribe events.""" - executor = copy(self) - executor.root_value = payload - executor.errors = None - return executor - - def execute_operation( - self, - serially: bool | None = None, - ) -> AwaitableOrValue[ExecutionResult | ExperimentalIncrementalExecutionResults]: - """Execute an operation. - - Implements the "Executing operations" section of the spec. - - Return a possible coroutine object that will eventually yield the data described - by the "Response" section of the GraphQL specification. - - If errors are encountered while executing a GraphQL field, only that field and - its descendants will be omitted, and sibling fields will still be executed. An - execution which encounters errors will still result in a coroutine object that - can be executed without errors. - - Errors from sub-fields of a NonNull type may propagate to the top level, - at which point we still log the error and null the parent field, which - in this case is the entire response. - - If the operation is aborted, the whole operation is rejected with an - aborted execution error rather than resolving to a partial response with - located errors; the partial result that the unwinding execution can still - produce is exposed on that error. - """ - abort_signal = self.abort_signal - if abort_signal is not None and abort_signal.aborted: - self.run_async_work_finished_hook() - raise self.abort_error() - try: - operation = self.operation - schema = self.schema - operation_type = operation.operation - root_type = schema.get_root_type(operation_type) - if root_type is None: - msg = ( - "Schema is not configured to execute" - f" {operation_type.value} operation." - ) - raise GraphQLError(msg, operation) # noqa: TRY301 - root_value = self.root_value - - collected_fields = collect_fields( - schema, - self.fragments, - self.variable_values, - root_type, - operation, - self.hide_suggestions, - ) - - grouped_field_set, new_defer_usages, _forbidden = collected_fields - - graphql_wrapped_result = ( - self.execute_execution_plan( - root_type, - root_value, - new_defer_usages, - build_execution_plan(grouped_field_set), - ) - if new_defer_usages - else self.execute_root_grouped_field_set( - root_type, - root_value, - grouped_field_set, - operation_type == OperationType.MUTATION - if serially is None - else serially, - None, - ) - ) - - if self.is_awaitable(graphql_wrapped_result): - - async def await_result() -> ( - ExecutionResult | ExperimentalIncrementalExecutionResults - ): - try: - resolved = await graphql_wrapped_result - except GraphQLError as error: - self.run_async_work_finished_hook() - return ExecutionResult(None, with_error(self.errors, error)) - except Exception: - # cancel incremental work started early and close the - # stream sources before re-raising, e.g. the abort reason - await self.cancel_incremental_work() - self.run_async_work_finished_hook() - raise - return self.build_data_response( - resolved.result, resolved.increments - ) - - if abort_signal is None: - return await_result() - return self.with_aborted_execution_error(await_result()) - - resolved = cast("GraphQLWrappedResult", graphql_wrapped_result) - - except GraphQLError as error: - self.run_async_work_finished_hook() - return self.finish(ExecutionResult(None, with_error(self.errors, error))) - - return self.finish( - self.build_data_response(resolved.result, resolved.increments) - ) - - def execute_execution_plan( - self, - return_type: GraphQLObjectType, - source_value: Any, - new_defer_usages: list[DeferUsage], - execution_plan: ExecutionPlan, - path: Path | None = None, - incremental_context: IncrementalContext | None = None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None = None, - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Execute an execution plan.""" - new_defer_map = get_new_defer_map(new_defer_usages, defer_map, path) - - grouped_field_set, new_grouped_field_sets = execution_plan - - graphql_wrapped_result = self.execute_fields( - return_type, - source_value, - path, - grouped_field_set, - incremental_context, - new_defer_map, - ) - - if new_grouped_field_sets: - new_pending_execution_groups = self.collect_execution_groups( - return_type, - source_value, - path, - incremental_context.defer_usage_set if incremental_context else None, - new_grouped_field_sets, - new_defer_map, - ) - return self.with_new_execution_groups( - graphql_wrapped_result, new_pending_execution_groups - ) - - return graphql_wrapped_result - - def execute_root_grouped_field_set( - self, - root_type: GraphQLObjectType, - root_value: Any, - grouped_field_set: GroupedFieldSet, - serially: bool, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Execute the root grouped field set.""" - return (self.execute_fields_serially if serially else self.execute_fields)( - root_type, - root_value, - None, - grouped_field_set, - None, - defer_map, - ) - - def execute_fields_serially( - self, - parent_type: GraphQLObjectType, - source_value: Any, - path: Path | None, - grouped_field_set: GroupedFieldSet, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Execute the given fields serially. - - Implements the "Executing selection sets" section of the spec - for fields that must be executed serially. - """ - is_awaitable = self.is_awaitable - abort_signal = self.abort_signal - - def reducer( - graphql_wrapped_result: GraphQLWrappedResult[dict[str, Any]], - field_item: tuple[str, FieldDetailsList], - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - response_name, field_details_list = field_item - field_path = Path(path, response_name, parent_type.name) - if abort_signal is not None and abort_signal.aborted: - # Reject the whole operation rather than serially completing the - # remaining fields with located abort errors. - raise self.abort_error() - result = self.execute_field( - parent_type, - source_value, - field_details_list, - field_path, - incremental_context, - defer_map, - ) - if result is Undefined: - return graphql_wrapped_result - if is_awaitable(result): - - async def set_result() -> GraphQLWrappedResult[dict[str, Any]]: - resolved = await result - graphql_wrapped_result.result[response_name] = resolved.result - graphql_wrapped_result.add_increments(resolved.increments) - return graphql_wrapped_result - - return set_result() - - resolved = cast("GraphQLWrappedResult[dict[str, Any]]", result) - graphql_wrapped_result.result[response_name] = resolved.result - graphql_wrapped_result.add_increments(resolved.increments) - return graphql_wrapped_result - - return async_reduce( - reducer, grouped_field_set.items(), GraphQLWrappedResult({}) - ) - - def execute_fields( - self, - parent_type: GraphQLObjectType, - source_value: Any, - path: Path | None, - grouped_field_set: GroupedFieldSet, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Execute the given fields concurrently. - - Implements the "Executing selection sets" section of the spec - for fields that may be executed in parallel. - """ - results: dict[str, Any] = {} - graphql_wrapped_result = GraphQLWrappedResult(results) - add_increments = graphql_wrapped_result.add_increments - is_awaitable = self.is_awaitable - awaitable_fields: list[str] = [] - append_awaitable = awaitable_fields.append - try: - for response_name, field_details_list in grouped_field_set.items(): - field_path = Path(path, response_name, parent_type.name) - result = self.execute_field( - parent_type, - source_value, - field_details_list, - field_path, - incremental_context, - defer_map, - ) - if result is not Undefined: - if is_awaitable(result): - - async def resolve( - result: Awaitable[GraphQLWrappedResult[dict[str, Any]]], - ) -> dict[str, Any]: - resolved = await result - add_increments(resolved.increments) - return resolved.result - - results[response_name] = resolve(result) - append_awaitable(response_name) - else: - result = cast("GraphQLWrappedResult[dict[str, Any]]", result) - results[response_name] = result.result - add_increments(result.increments) - except Exception: - if awaitable_fields: - # Ensure that awaitables created by other fields are settled, - # as they may also fail. - self.settle_in_background( - [results[field] for field in awaitable_fields] - ) - raise - - # If there are no coroutines, we can just return the object. - if not awaitable_fields: - return graphql_wrapped_result - - # Otherwise, results is a map from field name to the result of resolving that - # field, which is possibly a coroutine object. Return a coroutine object that - # will yield this same map, but with any coroutines awaited in parallel and - # replaced with the values they yielded. - async def get_results() -> GraphQLWrappedResult[dict[str, Any]]: - if len(awaitable_fields) == 1: - # If there is only one field, avoid the overhead of parallelization. - field = awaitable_fields[0] - results[field] = await results[field] - else: - awaited_results = await gather_with_cancel( - *(results[field] for field in awaitable_fields) - ) - results.update(zip(awaitable_fields, awaited_results, strict=True)) - - return GraphQLWrappedResult(results, graphql_wrapped_result.increments) - - return get_results() - - def execute_field( - self, - parent_type: GraphQLObjectType, - source: Any, - field_details_list: FieldDetailsList, - path: Path, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[Any]] | UndefinedType: - """Resolve the field on the given source object. - - Implements the "Executing fields" section of the spec. - - In particular, this method figures out the value that the field returns by - calling its resolve function, then calls complete_value to await coroutine - objects, coercing scalars, or execute the sub-selection-set for objects. - """ - first_field_details = field_details_list[0] - first_field_node = first_field_details.node - field_name = first_field_node.name.value - field_def = self.schema.get_field(parent_type, field_name) - if not field_def: - return Undefined - - return_type = field_def.type - resolve_fn = field_def.resolve or self.field_resolver - - if self.middleware_manager: - resolve_fn = self.middleware_manager.get_field_resolver(resolve_fn) - - info = self.build_resolve_info( - field_def, to_nodes(field_details_list), parent_type, path - ) - - # Get the resolve function, regardless of if its result is normal or abrupt - # (error). - try: - # Build a dictionary of arguments from the field.arguments AST, using the - # variables scope to fulfill any variable references. - args = get_argument_values( - field_def, - first_field_node, - self.variable_values, - first_field_details.fragment_variable_values, - self.hide_suggestions, - ) - - # Note that contrary to the JavaScript implementation, we pass the context - # value as part of the resolve info. - result = resolve_fn(source, info, **args) - - if self.is_awaitable(result): - return self.complete_awaitable_value( - return_type, - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - - completed = self.complete_value( - return_type, - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - if self.is_awaitable(completed): - - async def await_completed() -> Any: - try: - return await completed - except Exception as raw_error: - self.handle_field_error( - raw_error, - return_type, - field_details_list, - path, - incremental_context, - ) - return GraphQLWrappedResult(None) - - return await_completed() - - except Exception as raw_error: - self.handle_field_error( - raw_error, - return_type, - field_details_list, - path, - incremental_context, - ) - return GraphQLWrappedResult(None) - - return completed - - def build_resolve_info( - self, - field_def: GraphQLField, - field_nodes: list[FieldNode], - parent_type: GraphQLObjectType, - path: Path, - ) -> GraphQLResolveInfo: - """Build the GraphQLResolveInfo object. - - For internal use only. - """ - # The resolve function's first argument is a collection of information about - # the current execution state. - return GraphQLResolveInfo( - field_nodes[0].name.value, - field_nodes, - field_def.type, - parent_type, - path, - self.schema, - self.fragment_definitions, - self.root_value, - self.operation, - self.variable_values, - self.context_value, - self.is_awaitable, - self.abort_signal, - self.async_helpers, - ) - - def handle_field_error( - self, - raw_error: Exception, - return_type: GraphQLOutputType, - field_details_list: FieldDetailsList, - path: Path, - incremental_context: IncrementalContext | None = None, - ) -> None: - """Handle error properly according to the field type.""" - error = located_error(raw_error, to_nodes(field_details_list), path.as_list()) - - # If the field type is non-nullable, then it is resolved without any protection - # from errors, however it still properly locates the error. - if self.error_propagation and is_non_null_type(return_type): - raise error - - # Otherwise, error protection is applied, logging the error and resolving a - # null value for this field if one is encountered. - context = incremental_context or self - errors = context.errors - if errors is None: - context.errors = errors = [] - errors.append(error) - - def complete_value( - self, - return_type: GraphQLOutputType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - path: Path, - result: Any, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[Any]]: - """Complete a value. - - Implements the instructions for completeValue as defined in the - "Value completion" section of the spec. - - If the field type is Non-Null, then this recursively completes the value - for the inner type. It throws a field error if that completion returns null, - as per the "Nullability" section of the spec. - - If the field type is a List, then this recursively completes the value - for the inner type on each item in the list. - - If the field type is a Scalar or Enum, ensures the completed value is a legal - value of the type by calling the ``coerce_output_value`` method of GraphQL type - definition. - - If the field is an abstract type, determine the runtime type of the value and - then complete based on that type. - - Otherwise, the field type expects a sub-selection set, and will complete the - value by evaluating all sub-selections. - """ - # If result is an Exception, throw a located error. - if isinstance(result, Exception): - raise result - - # If field type is NonNull, complete for inner type, and throw field error if - # result is null. - if is_non_null_type(return_type): - completed = self.complete_value( - return_type.of_type, - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - if isinstance(completed, GraphQLWrappedResult) and completed.result is None: - msg = ( - "Cannot return null for non-nullable field" - f" {info.parent_type}.{info.field_name}." - ) - raise TypeError(msg) - return completed - - # If result value is null or undefined then return null. - if result is None or result is Undefined: - return GraphQLWrappedResult(None) - - # If field type is List, complete each item in the list with inner type - if is_list_type(return_type): - return self.complete_list_value( - return_type, - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - - # If field type is a leaf type, Scalar or Enum, coerce to a valid value, - # returning null if coercion is not possible. - if is_leaf_type(return_type): - return GraphQLWrappedResult(self.complete_leaf_value(return_type, result)) - - # If field type is an abstract type, Interface or Union, determine the runtime - # Object type and complete for that type. - if is_abstract_type(return_type): - return self.complete_abstract_value( - return_type, - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - - # If field type is Object, execute and complete all sub-selections. - if is_object_type(return_type): - return self.complete_object_value( - return_type, - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - - # Not reachable. All possible output types have been considered. - msg = ( - "Cannot complete value of unexpected output type:" - f" '{inspect(return_type)}'." - ) # pragma: no cover - raise TypeError(msg) # pragma: no cover - - async def with_abort_signal(self, awaitable: Awaitable[T]) -> T: - """Await a value, but cancel immediately if the abort signal is triggered. - - This wraps awaitables returned by resolvers (and awaitable list items) so - that a triggered abort signal interrupts execution *immediately* instead of - only at the next field boundary. Without this, a hanging asynchronous - resolver would prevent the operation from ever being cancelled. - - If the abort signal fires before the awaitable settles, the underlying - awaitable is cancelled and the abort reason is raised (an exception reason - is raised as is, any other value is reported as an unexpected error value). - """ - abort_signal = self.abort_signal - if abort_signal is None: - return await awaitable - task = ensure_future(awaitable) - if not abort_signal.aborted: - abort = ensure_future(abort_signal.wait()) - try: - await wait({task, abort}, return_when=FIRST_COMPLETED) - finally: - if not abort.done(): - abort.cancel() - if not abort_signal.aborted: - return task.result() - # The abort signal fired (possibly in the same tick the task settled); - # discard any task result and reject with the abort reason. - task.cancel() - with suppress(BaseException): - await task - raise self.abort_error() - - def abort_error(self) -> Exception: - """Return the exception to raise when execution has been aborted. - - An abort reason that is itself an exception is raised as is; any other value - is reported as an unexpected error value. - """ - reason = self.abort_signal.reason # type: ignore[union-attr] - if isinstance(reason, Exception): - return reason - msg = f"Unexpected error value: {inspect(reason)}" - return TypeError(msg) - - async def with_aborted_execution_error(self, awaitable: Awaitable[T]) -> T: - """Await a result, but raise an aborted execution error when aborted. - - Unlike :meth:`with_abort_signal`, the awaited work is not cancelled as a - whole when the abort signal is triggered: only its in-flight resolvers are - cancelled individually, so that it still settles into a partial result, - which is exposed via the raised aborted execution error. - """ - abort_signal = self.abort_signal - task = ensure_future(awaitable) - if not abort_signal.aborted: # type: ignore[union-attr] - abort = ensure_future(abort_signal.wait()) # type: ignore[union-attr] - try: - await wait({task, abort}, return_when=FIRST_COMPLETED) - finally: - if not abort.done(): - abort.cancel() - if not abort_signal.aborted: # type: ignore[union-attr] - return task.result() - # The abort signal fired (possibly in the same tick the task settled); - # let the partial result settle in the background and expose it on the - # aborted execution error that the operation is rejected with. - self.settle_in_background([task]) - raise self.create_aborted_execution_error(task) - - def finish(self, result: T) -> T: - """Check that execution has not been aborted before returning its result. - - If the operation has been aborted during otherwise synchronous execution, - raise an aborted execution error exposing the already completed result. - """ - abort_signal = self.abort_signal - if abort_signal is not None and abort_signal.aborted: - raise self.create_aborted_execution_error(result) - return result - - def create_aborted_execution_error( - self, result: AwaitableOrValue[Any] - ) -> AbortedGraphQLExecutionError: - """Create an aborted execution error exposing the given result.""" - reason = self.abort_signal.reason # type: ignore[union-attr] - return AbortedGraphQLExecutionError(reason, result) - - def box_incremental_result( - self, result: AwaitableOrValue[T] - ) -> BoxedAwaitableOrValue[T]: - """Box a possibly awaitable incremental result. - - A pending result is registered so that it can be cancelled when the - incremental execution is stopped before it has settled. - """ - boxed = BoxedAwaitableOrValue(result) - future = boxed.pending_future - if future is not None: - futures = self.pending_incremental_futures - futures.add(future) - future.add_done_callback(futures.discard) - return boxed - - async def cancel_incremental_work(self) -> None: - """Cancel all pending incremental work and close the stream sources. - - Cancels the still pending incremental execution tasks first and waits for - their cancellation to settle, so that no early execution continues and no - iteration is pending on the stream sources any more, then triggers and - awaits the early return of all remaining cancellable streams. - """ - futures = self.pending_incremental_futures - if futures: - pending = list(futures) - for future in pending: - future.cancel() - await gather(*pending, return_exceptions=True) - cancellable_streams = self.cancellable_streams - if cancellable_streams: - early_returns = [ - early_return - for early_return in ( - stream_record.early_return() - for stream_record in cancellable_streams - ) - if default_is_awaitable(early_return) - ] - cancellable_streams.clear() - if early_returns: - await gather(*early_returns, return_exceptions=True) - - def settle_in_background(self, awaitables: list[Awaitable[Any]]) -> None: - """Settle the given pending awaitables in the background. - - A bubbling synchronous error must not wait for pending sibling awaitables, - but they must still be settled so that their errors can be observed before - they would be orphaned (the Python analog of silencing JS unhandled - rejections). Without a running event loop the awaitables can never run; - pending coroutines are then closed instead to dispose of them. - """ - try: - get_running_loop() - except RuntimeError: # no running event loop - for awaitable in awaitables: - if iscoroutine(awaitable): # pragma: no branch - awaitable.close() - return - future = gather(*awaitables, return_exceptions=True) - background_futures = self.background_futures - background_futures.add(future) - future.add_done_callback(background_futures.discard) - - def track_async_work(self, values: Sequence[Any]) -> None: - """Track possibly awaitable values as pending asynchronous work. - - Awaitables among the given values are settled in the background, so that - they are still settled and their errors observed when they would otherwise - be abandoned. Non-awaitable values are ignored. - """ - is_awaitable = self.is_awaitable - awaitables: list[Awaitable[Any]] = [ - value for value in values if is_awaitable(value) - ] - if awaitables: - self.settle_in_background(awaitables) - - def gather_async_work( - self, values: Sequence[Awaitable[Any]] - ) -> Awaitable[list[Any]]: - """Concurrently await the given values as one unit of asynchronous work. - - This allows resolvers to await multiple concurrent operations together. - When one of the values fails, the others are cancelled and settled before - the error is propagated, so that no asynchronous work is orphaned. - """ - return gather_with_cancel(*values) - - def run_async_work_finished_hook(self) -> None: - """Run the hook signaling that all asynchronous work has finished. - - If an ``async_work_finished`` execution hook is provided, run it as soon as - all tracked pending asynchronous work has been settled - synchronously when - there is no pending asynchronous work, which allows synchronous execution - paths to remain synchronous. Errors raised by the hook are ignored. - """ - hooks = self.hooks - hook = hooks.async_work_finished if hooks is not None else None - if hook is None: - return - info = AsyncWorkFinishedInfo(self) - background_futures = self.background_futures - if not background_futures: - with suppress_exceptions: - hook(info) - return - - async def wait_and_run_hook() -> None: - while background_futures: - await wait(list(background_futures)) - with suppress_exceptions: - hook(info) - - # keep a reference to the task so that it is not garbage collected - self.async_work_finished_hook_task = ensure_future(wait_and_run_hook()) - - def cancellable_iterable(self, iterable: AsyncIterable[T]) -> AsyncIterable[T]: - """Wrap an async iterable so pending iteration is cancelled on abort. - - When the abort signal is triggered, any pending ``__anext__`` call returns - immediately by raising the abort reason. This mirrors the JavaScript - ``cancellableIterable``; GraphQL-core needs no ``AbortSignalListener`` - class since :meth:`with_abort_signal` already provides the cancellation - mechanism. - """ - if self.abort_signal is None: - return iterable - with_abort_signal = self.with_abort_signal - iterator = iterable.__aiter__() - - class CancellableAsyncIterator: - def __aiter__(self) -> AsyncIterator[T]: - return self - - def __anext__(self) -> Awaitable[T]: - return with_abort_signal(iterator.__anext__()) - - async def aclose(self) -> None: - aclose = getattr(iterator, "aclose", None) - if aclose is not None: - await aclose() - - return CancellableAsyncIterator() - - async def complete_awaitable_value( - self, - return_type: GraphQLOutputType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - path: Path, - result: Any, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> GraphQLWrappedResult[Any]: - """Complete an awaitable value.""" - try: - resolved = await self.with_abort_signal(result) - completed = self.complete_value( - return_type, - field_details_list, - info, - path, - resolved, - incremental_context, - defer_map, - ) - if self.is_awaitable(completed): - completed = await completed - except Exception as raw_error: - self.handle_field_error( - raw_error, return_type, field_details_list, path, incremental_context - ) - completed = GraphQLWrappedResult(None) - return completed # type: ignore - - def get_stream_usage( - self, field_details_list: FieldDetailsList, path: Path - ) -> StreamUsage | None: - """Get stream usage. - - Returns an object containing info for streaming if a field should be - streamed based on the experimental flag, stream directive present and - not disabled by the "if" argument. - """ - # do not stream inner lists of multidimensional lists - if isinstance(path.key, int): - return None - - stream_usage = self._stream_usages.get(field_details_list) - if stream_usage is not None: - return stream_usage # pragma: no cover - - # validation only allows equivalent streams on multiple fields, so it is - # safe to only check the first field_node for the stream directive - stream = get_directive_values( - GraphQLStreamDirective, field_details_list[0].node, self.variable_values - ) - - if not stream or stream.get("if") is False: - return None - - initial_count = stream.get("initialCount") - if initial_count is None or initial_count < 0: - msg = "initialCount must be a positive integer" - raise ValueError(msg) - - if self.operation.operation == OperationType.SUBSCRIPTION: - msg = ( - "`@stream` directive not supported on subscription operations." - " Disable `@stream` by setting the `if` argument to `false`." - ) - raise TypeError(msg) - - streamed_field_details_list: FieldDetailsList = [ - FieldDetails(field_details.node, None) - for field_details in field_details_list - ] - - stream_usage = StreamUsage( - stream.get("label"), stream["initialCount"], streamed_field_details_list - ) - - self._stream_usages[field_details_list] = stream_usage - - return stream_usage - - async def complete_async_iterator_value( - self, - item_type: GraphQLOutputType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - path: Path, - async_iterator: AsyncIterator[Any], - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> GraphQLWrappedResult[list[Any]]: - """Complete an async iterator. - - Complete an async iterator value by completing the result and calling - recursively until all the results are completed. - """ - is_awaitable = self.is_awaitable - complete_list_item_value = self.complete_list_item_value - complete_awaitable_list_item_value = self.complete_awaitable_list_item_value - completed_results: list[Any] = [] - append_completed = completed_results.append - graphql_wrapped_result = GraphQLWrappedResult(completed_results) - add_increment = graphql_wrapped_result.add_increment - awaitable_indices: list[int] = [] - append_awaitable = awaitable_indices.append - stream_usage = self.get_stream_usage(field_details_list, path) - try: - early_return = async_iterator.aclose # type: ignore[attr-defined] - except AttributeError: - early_return = None - index = 0 - try: - while True: - if stream_usage and index >= stream_usage.initial_count: - stream_item_queue = self.build_async_stream_item_queue( - index, - path, - async_iterator, - stream_usage.field_details_list, - info, - item_type, - ) - - stream_record: StreamRecord - - if early_return is None: - stream_record = StreamRecord( - stream_item_queue, path, stream_usage.label - ) - else: - stream_record = CancellableStreamRecord( - early_return, - stream_item_queue, - path, - stream_usage.label, - ) - if self.cancellable_streams is None: # pragma: no branch - self.cancellable_streams = set() - self.cancellable_streams.add(stream_record) - - add_increment(stream_record) - break - - item_path = path.add_key(index, None) - try: - item = await anext(async_iterator) - except StopAsyncIteration: - break - except Exception as raw_error: - raise located_error( - raw_error, to_nodes(field_details_list), path.as_list() - ) from raw_error - - if is_awaitable(item): - append_completed( - complete_awaitable_list_item_value( - item, - graphql_wrapped_result, - item_type, - field_details_list, - info, - item_path, - incremental_context, - defer_map, - ) - ) - append_awaitable(index) - - elif complete_list_item_value( - item, - completed_results, - graphql_wrapped_result, - item_type, - field_details_list, - info, - item_path, - incremental_context, - defer_map, - ): - append_awaitable(index) - - index += 1 - except Exception: - if early_return is not None: # pragma: no branch - with suppress_exceptions: - await early_return() - if awaitable_indices: - # Settle any awaitable items already collected in the background, - # so that the current error is not delayed. - self.settle_in_background( - [completed_results[index] for index in awaitable_indices] - ) - raise - - if not awaitable_indices: - return graphql_wrapped_result - - if len(awaitable_indices) == 1: - # If there is only one index, avoid the overhead of parallelization. - index = awaitable_indices[0] - completed_results[index] = await completed_results[index] - else: - awaited_results = await gather_with_cancel( - *(completed_results[index] for index in awaitable_indices) - ) - for index, sub_result in zip( - awaitable_indices, awaited_results, strict=True - ): - completed_results[index] = sub_result - return GraphQLWrappedResult( - completed_results, graphql_wrapped_result.increments - ) - - def complete_list_value( - self, - return_type: GraphQLList[GraphQLOutputType], - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - path: Path, - result: AsyncIterable[Any] | Iterable[Any], - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[list[Any]]]: - """Complete a list value. - - Complete a list value by completing each item in the list with the inner type. - """ - item_type = return_type.of_type - - if self.is_async_iterable(result): - async_iterator = self.cancellable_iterable(result).__aiter__() - - return self.complete_async_iterator_value( - item_type, - field_details_list, - info, - path, - async_iterator, - incremental_context, - defer_map, - ) - - if not is_iterable(result): - msg = ( - "Expected Iterable, but did not find one for field" - f" '{info.parent_type}.{info.field_name}'." - ) - raise GraphQLError(msg) - - return self.complete_iterable_value( - item_type, - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - - def complete_iterable_value( - self, - item_type: GraphQLOutputType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - path: Path, - items: Iterable[Any], - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[list[Any]]]: - """Complete an iterable value.""" - # This is specified as a simple map, however we're optimizing the path - # where the list contains no awaitable routine objects by avoiding creating - # another awaitable object. - is_awaitable = self.is_awaitable - complete_list_item_value = self.complete_list_item_value - complete_awaitable_list_item_value = self.complete_awaitable_list_item_value - completed_results: list[Any] = [] - graphql_wrapped_result = GraphQLWrappedResult(completed_results) - add_increment = graphql_wrapped_result.add_increment - append_completed = completed_results.append - awaitable_indices: list[int] = [] - append_awaitable = awaitable_indices.append - stream_usage = self.get_stream_usage(field_details_list, path) - iterator = iter(items) - index = 0 - try: - while True: - try: - item = next(iterator) - except StopIteration: - break - if stream_usage and index >= stream_usage.initial_count: - sync_stream_item_queue = self.build_sync_stream_item_queue( - item, - index, - path, - iterator, - stream_usage.field_details_list, - info, - item_type, - ) - sync_stream_record = StreamRecord( - sync_stream_item_queue, path, stream_usage.label - ) - - add_increment(sync_stream_record) - break - - # No need to modify the info object containing the path, - # since from here on it is not ever accessed by resolver functions. - item_path = path.add_key(index, None) - - if is_awaitable(item): - append_completed( - complete_awaitable_list_item_value( - item, - graphql_wrapped_result, - item_type, - field_details_list, - info, - item_path, - incremental_context, - defer_map, - ) - ) - append_awaitable(index) - - elif complete_list_item_value( - item, - completed_results, - graphql_wrapped_result, - item_type, - field_details_list, - info, - item_path, - incremental_context, - defer_map, - ): - append_awaitable(index) - - index += 1 - except Exception: - # Do not close the iterator. Instead, drain it so that any awaitable - # items it still holds can be settled in the background before they - # would be orphaned. - maybe_awaitables = [completed_results[index] for index in awaitable_indices] - maybe_awaitables.extend(collect_iterator_awaitables(iterator, is_awaitable)) - if maybe_awaitables: - self.settle_in_background(maybe_awaitables) - raise - - if not awaitable_indices: - return graphql_wrapped_result - - async def get_completed_results() -> GraphQLWrappedResult[list[Any]]: - if len(awaitable_indices) == 1: - # If there is only one index, avoid the overhead of parallelization. - index = awaitable_indices[0] - completed_results[index] = await completed_results[index] - else: - awaited_results = await gather_with_cancel( - *(completed_results[index] for index in awaitable_indices) - ) - for index, sub_result in zip( - awaitable_indices, awaited_results, strict=True - ): - completed_results[index] = sub_result - return GraphQLWrappedResult( - completed_results, graphql_wrapped_result.increments - ) - - return get_completed_results() - - def complete_list_item_value( - self, - item: Any, - complete_results: list[Any], - parent: GraphQLWrappedResult[list[Any]], - item_type: GraphQLOutputType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - item_path: Path, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> bool: - """Complete a list item value by adding it to the completed results. - - Returns True if the value is awaitable. - """ - is_awaitable = self.is_awaitable - - try: - completed_item = self.complete_value( - item_type, - field_details_list, - info, - item_path, - item, - incremental_context, - defer_map, - ) - - if is_awaitable(completed_item): - - async def await_completed() -> Any: - try: - resolved = await completed_item # type: ignore - except Exception as raw_error: - self.handle_field_error( - raw_error, - item_type, - field_details_list, - item_path, - incremental_context, - ) - return None - parent.add_increments(resolved.increments) - return resolved.result - - complete_results.append(await_completed()) - return True - - completed_item = cast("GraphQLWrappedResult[Any]", completed_item) - complete_results.append(completed_item.result) - parent.add_increments(completed_item.increments) - - except Exception as raw_error: - self.handle_field_error( - raw_error, - item_type, - field_details_list, - item_path, - incremental_context, - ) - complete_results.append(None) - - return False - - async def complete_awaitable_list_item_value( - self, - item: Any, - parent: GraphQLWrappedResult[list[Any]], - item_type: GraphQLOutputType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - item_path: Path, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> Any: - """Complete an awaitable list item value.""" - try: - resolved = await self.with_abort_signal(item) - completed = self.complete_value( - item_type, - field_details_list, - info, - item_path, - resolved, - incremental_context, - defer_map, - ) - if self.is_awaitable(completed): - completed = await completed - completed = cast("GraphQLWrappedResult[list[Any]]", completed) - parent.add_increments(completed.increments) - except Exception as raw_error: - self.handle_field_error( - raw_error, - item_type, - field_details_list, - item_path, - incremental_context, - ) - return None - else: - return completed.result - - @staticmethod - def complete_leaf_value(return_type: GraphQLLeafType, result: Any) -> Any: - """Complete a leaf value. - - Complete a Scalar or Enum by coercing to a valid value, returning null if - coercion is not possible. - """ - coerced = return_type.coerce_output_value(result) - if coerced is Undefined or coerced is None: - msg = ( - f"Expected `{inspect(return_type)}.coerce_output_value(" - f"{inspect(result)})` to return non-nullable value, returned:" - f" {inspect(coerced)}" - ) - raise TypeError(msg) - return coerced - - def complete_abstract_value( - self, - return_type: GraphQLAbstractType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - path: Path, - result: Any, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Complete an abstract value. - - Complete a value of an abstract type by determining the runtime object type of - that value, then complete the value for that type. - """ - resolve_type_fn = return_type.resolve_type or self.type_resolver - runtime_type = resolve_type_fn(result, info, return_type) - - if self.is_awaitable(runtime_type): - - async def await_complete_object_value() -> Any: - value = self.complete_object_value( - self.ensure_valid_runtime_type( - await self.with_abort_signal(runtime_type), # type: ignore - return_type, - field_details_list, - info, - result, - ), - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - if self.is_awaitable(value): - return await value - return value # pragma: no cover - - return await_complete_object_value() - runtime_type = cast("str | None", runtime_type) - - return self.complete_object_value( - self.ensure_valid_runtime_type( - runtime_type, return_type, field_details_list, info, result - ), - field_details_list, - info, - path, - result, - incremental_context, - defer_map, - ) - - def ensure_valid_runtime_type( - self, - runtime_type_name: Any, - return_type: GraphQLAbstractType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - result: Any, - ) -> GraphQLObjectType: - """Ensure that the given type is valid at runtime.""" - if runtime_type_name is None: - msg = ( - f"Abstract type '{return_type}' must resolve" - " to an Object type at runtime" - f" for field '{info.parent_type}.{info.field_name}'." - f" Either the '{return_type}' type should provide" - " a 'resolve_type' function or each possible type should provide" - " an 'is_type_of' function." - ) - raise GraphQLError(msg, to_nodes(field_details_list)) - - if not isinstance(runtime_type_name, str): - msg = ( - f"Abstract type '{return_type}' must resolve" - " to an Object type at runtime" - f" for field '{info.parent_type}.{info.field_name}' with value" - f" {inspect(result)}, received '{inspect(runtime_type_name)}'," - " which is not a valid Object type name." - ) - raise GraphQLError(msg, to_nodes(field_details_list)) - - runtime_type = self.schema.get_type(runtime_type_name) - - if runtime_type is None: - msg = ( - f"Abstract type '{return_type}' was resolved to a type" - f" '{runtime_type_name}' that does not exist inside the schema." - ) - raise GraphQLError(msg, to_nodes(field_details_list)) - - if not is_object_type(runtime_type): - msg = ( - f"Abstract type '{return_type}' was resolved" - f" to a non-object type '{runtime_type_name}'." - ) - raise GraphQLError(msg, to_nodes(field_details_list)) - - if not self.schema.is_sub_type(return_type, runtime_type): - msg = ( - f"Runtime Object type '{runtime_type}' is not a possible" - f" type for '{return_type}'." - ) - raise GraphQLError(msg, to_nodes(field_details_list)) - - return runtime_type - - def complete_object_value( - self, - return_type: GraphQLObjectType, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - path: Path, - result: Any, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Complete an Object value by executing all sub-selections.""" - # If there is an `is_type_of()` predicate function, call it with the current - # result. If `is_type_of()` returns False, then raise an error rather than - # continuing execution. - if return_type.is_type_of: - is_type_of = return_type.is_type_of(result, info) - - if self.is_awaitable(is_type_of): - - async def execute_subfields_async() -> GraphQLWrappedResult[ - dict[str, Any] - ]: - if not await self.with_abort_signal(is_type_of): - raise invalid_return_type_error( - return_type, result, field_details_list - ) - graphql_wrapped_result = self.collect_and_execute_subfields( - return_type, - field_details_list, - path, - result, - incremental_context, - defer_map, - ) - if self.is_awaitable(graphql_wrapped_result): # pragma: no cover - return await graphql_wrapped_result - return cast( - "GraphQLWrappedResult[dict[str, Any]]", graphql_wrapped_result - ) - - return execute_subfields_async() - - if not is_type_of: - raise invalid_return_type_error(return_type, result, field_details_list) - - return self.collect_and_execute_subfields( - return_type, - field_details_list, - path, - result, - incremental_context, - defer_map, - ) - - def collect_and_execute_subfields( - self, - return_type: GraphQLObjectType, - field_details_list: FieldDetailsList, - path: Path, - result: Any, - incremental_context: IncrementalContext | None, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Collect sub-fields to execute to complete this value.""" - collected_subfields = self.collect_subfields(return_type, field_details_list) - grouped_field_set, new_defer_usages, _forbidden = collected_subfields - - return ( - self.execute_fields( - return_type, result, path, grouped_field_set, incremental_context, None - ) - if defer_map is None and not new_defer_usages - else self.execute_execution_plan( - return_type, - result, - new_defer_usages, - self.build_sub_execution_plan( - grouped_field_set, - incremental_context.defer_usage_set - if incremental_context - else None, - ), - path, - incremental_context, - defer_map, - ) - ) - - def collect_subfields( - self, return_type: GraphQLObjectType, field_details_list: FieldDetailsList - ) -> CollectedFields: - """Collect subfields. - - A memoized function collecting relevant subfields regarding the return type. - Memoizing ensures the subfields are not repeatedly calculated, which saves - overhead when resolving lists of values. - """ - relevant_sub_fields = self._relevant_sub_fields - # We cannot use the field_details_list itself as key for the cache, since it - # is not hashable as a list. We also do not want to use the field_details_list - # itself (converted to a tuple) as keys, since hashing them is slow. - # Therefore, we use the ids of the field_details_list items as keys. Note that - # we do not use the id of the list, since we want to hit the cache for all - # lists of the same nodes, not only for the same list of nodes. Also, the list - # id may even be reused, in which case we would get wrong results from cache. - key = ( - (return_type, id(field_details_list[0])) - if len(field_details_list) == 1 # optimize most frequent case - else (return_type, *map(id, field_details_list)) - ) - collected_fields: CollectedFields | None = relevant_sub_fields.get(key) - if collected_fields is None: - collected_fields = collect_subfields( - self.schema, - self.fragments, - self.variable_values, - self.operation, - return_type, - field_details_list, - self.hide_suggestions, - ) - relevant_sub_fields[key] = collected_fields - return collected_fields - - def build_sub_execution_plan( - self, - original_grouped_field_set: GroupedFieldSet, - defer_usage_set: DeferUsageSet | None, - ) -> ExecutionPlan: - """Build a cached sub-execution plan.""" - execution_plans = self._execution_plans - exeecution_plan = execution_plans.get(original_grouped_field_set) - if exeecution_plan is not None: - return exeecution_plan - exeecution_plan = build_execution_plan( - original_grouped_field_set, defer_usage_set - ) - execution_plans[original_grouped_field_set] = exeecution_plan - return exeecution_plan - - def collect_execution_groups( - self, - parent_type: GraphQLObjectType, - source_value: Any, - path: Path | None, - parent_defer_usages: DeferUsageSet | None, - new_grouped_field_sets: RefMap[DeferUsageSet, GroupedFieldSet], - defer_map: RefMap[DeferUsage, DeferredFragmentRecord], - ) -> list[PendingExecutionGroup]: - """Execute deferred grouped field sets.""" - is_awaitable = self.is_awaitable - new_pending_execution_groups: list[PendingExecutionGroup] = [] - append_record = new_pending_execution_groups.append - for defer_usage_set, grouped_field_set in new_grouped_field_sets.items(): - deferred_fragment_records = get_deferred_fragment_records( - defer_usage_set, defer_map - ) - - pending_group = PendingExecutionGroup( - deferred_fragment_records, - cast("BoxedAwaitableOrValue[CompletedExecutionGroup]", None), - ) - - def executor( - pending_group: PendingExecutionGroup = pending_group, - grouped_field_set: GroupedFieldSet = grouped_field_set, - defer_usage_set: DeferUsageSet = defer_usage_set, - ) -> AwaitableOrValue[CompletedExecutionGroup]: - return self.execute_execution_group( - pending_group, - parent_type, - source_value, - path, - grouped_field_set, - IncrementalContext(defer_usage_set), - defer_map, - ) - - if self.enable_early_execution: - if should_defer(parent_defer_usages, defer_usage_set): - - async def execute_async( - executor: Callable[ - [], AwaitableOrValue[CompletedExecutionGroup] - ] = executor, - ) -> CompletedExecutionGroup: - result = executor() - if is_awaitable(result): - return await result - return result # type: ignore - - pending_group.result = self.box_incremental_result(execute_async()) - else: - pending_group.result = self.box_incremental_result(executor()) - else: - - def execute_sync( - executor: Callable[ - [], AwaitableOrValue[CompletedExecutionGroup] - ] = executor, - ) -> BoxedAwaitableOrValue[CompletedExecutionGroup]: - return self.box_incremental_result(executor()) - - pending_group.result = execute_sync - - append_record(pending_group) - - return new_pending_execution_groups - - def execute_execution_group( - self, - pending_execution_group: PendingExecutionGroup, - parent_type: GraphQLObjectType, - source_value: Any, - path: Path | None, - grouped_field_set: GroupedFieldSet, - incremental_context: IncrementalContext, - defer_map: RefMap[DeferUsage, DeferredFragmentRecord], - ) -> AwaitableOrValue[CompletedExecutionGroup]: - """Execute deferred grouped field set.""" - try: - result = self.execute_fields( - parent_type, - source_value, - path, - grouped_field_set, - incremental_context, - defer_map, - ) - except GraphQLError as error: - return FailedExecutionGroup( - pending_execution_group, - path.as_list() if path else [], - with_error(incremental_context.errors, error), - ) - - if self.is_awaitable(result): - - async def await_result() -> CompletedExecutionGroup: - try: - awaited_result = await result - except GraphQLError as error: - return FailedExecutionGroup( - pending_execution_group, - path.as_list() if path else [], - with_error(incremental_context.errors, error), - ) - return build_completed_execution_group( - incremental_context.errors, - pending_execution_group, - path, - awaited_result, - ) - - return await_result() - - return build_completed_execution_group( - incremental_context.errors, - pending_execution_group, - path, - result, # type: ignore - ) - - def build_sync_stream_item_queue( - self, - initial_item: AwaitableOrValue[Any], - initial_index: int, - stream_path: Path, - iterator: Iterable[Any], - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - item_type: GraphQLOutputType, - ) -> list[StreamItemRecord]: - """Build sync stream item queue.""" - is_awaitable = self.is_awaitable - enable_early_execution = self.enable_early_execution - complete_stream_item = self.complete_stream_item - - stream_item_queue: list[StreamItemRecord] = [] - append_stream_item = stream_item_queue.append - - def first_executor() -> StreamItemResult: - initial_path = stream_path.add_key(initial_index) - - first_stream_item: BoxedAwaitableOrValue[StreamItemResult] = ( - self.box_incremental_result( - complete_stream_item( - initial_path, - initial_item, - IncrementalContext(), - field_details_list, - info, - item_type, - ) - ) - ) - - current_index = initial_index + 1 - current_stream_item: ( - BoxedAwaitableOrValue[StreamItemResult] - | Callable[[], BoxedAwaitableOrValue[StreamItemResult]] - ) = first_stream_item - for item in iterator: - if isinstance(current_stream_item, BoxedAwaitableOrValue): - result = current_stream_item.value - if not is_awaitable(result) and result.errors: - break - - item_path = stream_path.add_key(current_index) - - def current_executor( - item: Any = item, item_path: Path = item_path - ) -> AwaitableOrValue[StreamItemResult]: - return complete_stream_item( - item_path, - item, - IncrementalContext(), - field_details_list, - info, - item_type, - ) - - current_stream_item = ( - self.box_incremental_result(current_executor()) - if enable_early_execution - else lambda executor=current_executor: self.box_incremental_result( - executor() - ) - ) - - append_stream_item(current_stream_item) - - current_index = initial_index + 1 - - append_stream_item(BoxedAwaitableOrValue(StreamItemResult())) - - return first_stream_item.value - - if enable_early_execution: - - async def await_first_stream_item() -> StreamItemResult: - return first_executor() - - append_stream_item(self.box_incremental_result(await_first_stream_item())) - else: - append_stream_item(lambda: self.box_incremental_result(first_executor())) - - return stream_item_queue - - def build_async_stream_item_queue( - self, - initial_index: int, - stream_path: Path, - async_iterator: AsyncIterator[Any], - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - item_type: GraphQLOutputType, - ) -> list[StreamItemRecord]: - """Build async stream item queue.""" - - def executor() -> AwaitableOrValue[StreamItemResult]: - return self.get_next_async_stream_item_result( - stream_item_queue, - stream_path, - initial_index, - async_iterator, - field_details_list, - info, - item_type, - ) - - stream_item_queue: list[StreamItemRecord] = [] - stream_item_queue.append( - self.box_incremental_result(executor()) - if self.enable_early_execution - else lambda: self.box_incremental_result(executor()) - ) - - return stream_item_queue - - async def get_next_async_stream_item_result( - self, - stream_item_queue: list[StreamItemRecord], - stream_path: Path, - index: int, - async_iterator: AsyncIterator[Any], - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - item_type: GraphQLOutputType, - ) -> StreamItemResult: - """Get the next async stream items result.""" - try: - item = await anext(async_iterator) - except StopAsyncIteration: - return StreamItemResult() - except Exception as error: - return StreamItemResult( - errors=[ - located_error( - error, to_nodes(field_details_list), stream_path.as_list() - ) - ], - ) - - item_path = stream_path.add_key(index) - - result = self.complete_stream_item( - item_path, - item, - IncrementalContext(), - field_details_list, - info, - item_type, - ) - - def executor() -> AwaitableOrValue[StreamItemResult]: - return self.get_next_async_stream_item_result( - stream_item_queue, - stream_path, - index + 1, - async_iterator, - field_details_list, - info, - item_type, - ) - - stream_item_queue.append( - self.box_incremental_result(executor()) - if self.enable_early_execution - else lambda: self.box_incremental_result(executor()) - ) - - if self.is_awaitable(result): - await sleep(0) # allow other tasks to run - return await result - return cast("StreamItemResult", result) - - def complete_stream_item( - self, - item_path: Path, - item: Any, - incremental_context: IncrementalContext, - field_details_list: FieldDetailsList, - info: GraphQLResolveInfo, - item_type: GraphQLOutputType, - ) -> AwaitableOrValue[StreamItemResult]: - """Complete the stream items.""" - is_awaitable = self.is_awaitable - if is_awaitable(item): - - async def await_stream_item_result() -> StreamItemResult: - try: - awaited_item = await self.complete_awaitable_value( - item_type, - field_details_list, - info, - item_path, - item, - incremental_context, - RefMap(), - ) - except GraphQLError as error: - return StreamItemResult( - errors=with_error(incremental_context.errors, error) - ) - return build_stream_item_result( - awaited_item, incremental_context.errors - ) - - return await_stream_item_result() - - try: - try: - result = self.complete_value( - item_type, - field_details_list, - info, - item_path, - item, - incremental_context, - RefMap(), - ) - except Exception as raw_error: - self.handle_field_error( - raw_error, - item_type, - field_details_list, - item_path, - incremental_context, - ) - result = GraphQLWrappedResult(None) - except GraphQLError as error: - return StreamItemResult( - errors=with_error(incremental_context.errors, error) - ) - - if is_awaitable(result): - - async def await_stream_item_result() -> StreamItemResult: - try: - try: - awaited_item = await result - except Exception as raw_error: - self.handle_field_error( - raw_error, - item_type, - field_details_list, - item_path, - incremental_context, - ) - awaited_item = GraphQLWrappedResult(None) - except GraphQLError as error: - return StreamItemResult( - errors=with_error(incremental_context.errors, error) - ) - return build_stream_item_result( - awaited_item, incremental_context.errors - ) - - return await_stream_item_result() - - return build_stream_item_result( - result, # type: ignore - incremental_context.errors, - ) - - def with_new_execution_groups( - self, - result: AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]], - new_pending_execution_groups: list[PendingExecutionGroup], - ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]: - """Add new pending execution groups to result.""" - if self.is_awaitable(result): - - async def await_result() -> GraphQLWrappedResult[dict[str, Any]]: - resolved = await result - resolved.add_increments(new_pending_execution_groups) - return resolved - - return await_result() - - resolved = cast("GraphQLWrappedResult[dict[str, Any]]", result) - resolved.add_increments(new_pending_execution_groups) - return resolved - - def build_data_response( - self, - data: dict[str, Any], - incremental_data_records: Sequence[IncrementalDataRecord] | None, - ) -> ExecutionResult | ExperimentalIncrementalExecutionResults: - """Build the data response.""" - if not incremental_data_records: - self.run_async_work_finished_hook() - return ExecutionResult(data, self.errors or None) - return build_incremental_response( - self, - data, - self.errors, - incremental_data_records, - ) - - -def with_error( - errors: Sequence[GraphQLError] | None, error: GraphQLError -) -> list[GraphQLError]: - """Return a new list of errors with the given error appended.""" - return [error] if errors is None else [*errors, error] - - -def to_nodes(field_details_list: FieldDetailsList) -> list[FieldNode]: - """Convert a field group to a list of field nodes.""" - return [field_details.node for field_details in field_details_list] - - -T = TypeVar("T") - - -class GraphQLWrappedResult(Generic[T]): - """Wrapper class for GraphQL results with increments.""" - - __slots__ = "increments", "result" - - result: T - increments: list[IncrementalDataRecord] | None - - def __init__( - self, - result: T, - increments: list[IncrementalDataRecord] | None = None, - ) -> None: - self.result = result - self.increments = increments - - def add_increment( - self, - increment: IncrementalDataRecord, - ) -> None: - """Add a single given increment to the wrapped result.""" - if self.increments is None: - self.increments = [increment] - else: - self.increments.append(increment) - - def add_increments( - self, - increments: Sequence[IncrementalDataRecord] | None, - ) -> None: - """Add the given increments to the wrapped result.""" - if increments is None: - return - if self.increments is None: - self.increments = list(increments) - else: - self.increments.extend(increments) - - UNEXPECTED_EXPERIMENTAL_DIRECTIVES = ( "The provided schema unexpectedly contains experimental directives" " (@defer or @stream). These directives may only be utilized" @@ -2461,12 +96,6 @@ def add_increments( ) -UNEXPECTED_MULTIPLE_PAYLOADS = ( - "Executing this GraphQL operation would unexpectedly produce multiple payloads" - " (due to @defer or @stream directive)" -) - - def execute( # noqa: PLR0913 schema: GraphQLSchema, document: DocumentNode, @@ -2573,7 +202,7 @@ def experimental_execute_incrementally( # noqa: PLR0913 and a stream of `subsequent_results`. """ if executor_class is None: - executor_class = Executor + executor_class = IncrementalExecutor # If a valid executor cannot be created due to incorrect arguments, # a "Response" with only errors is returned. @@ -2674,209 +303,6 @@ def execute_sync( return cast("ExecutionResult", result) -def invalid_return_type_error( - return_type: GraphQLObjectType, result: Any, field_details_list: FieldDetailsList -) -> GraphQLError: - """Create a GraphQLError for an invalid return type.""" - return GraphQLError( - f"Expected value of type '{return_type}' but got: {inspect(result)}.", - to_nodes(field_details_list), - ) - - -def deferred_fragment_record_from_defer_usage( - defer_usage: DeferUsage, defer_map: RefMap[DeferUsage, DeferredFragmentRecord] -) -> DeferredFragmentRecord: - """Get the deferred fragment record mapped to the given defer usage.""" - return defer_map[defer_usage] - - -def get_new_defer_map( - new_defer_usages: Sequence[DeferUsage], - defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None, - path: Path | None = None, -) -> RefMap[DeferUsage, DeferredFragmentRecord]: - """Get the new defer map. - - Instantiates new DeferredFragmentRecords for the given path within an - incremental data record, returning an updated map of DeferUsage - objects to DeferredFragmentRecords. - - Note: As defer directives may be used with operations returning lists, - a DeferUsage object may correspond to many DeferredFragmentRecords. - """ - new_defer_map = RefMap(None if defer_map is None else defer_map.items()) - - # For each new DeferUsage object: - for new_defer_usage in new_defer_usages: - parent_defer_usage = new_defer_usage.parent_defer_usage - - parent = ( - None - if parent_defer_usage is None - else deferred_fragment_record_from_defer_usage( - parent_defer_usage, new_defer_map - ) - ) - - # Instantiate the new record. - deferred_fragment_record = DeferredFragmentRecord( - path, new_defer_usage.label, parent - ) - - # Update the map. - new_defer_map[new_defer_usage] = deferred_fragment_record - - return new_defer_map - - -def should_defer( - parent_defer_usages: DeferUsageSet | None, defer_usage_set: DeferUsageSet -) -> bool: - """Decide whether to defer the given defer usage set. - - If we have a new child defer usage, defer. - Otherwise, this defer usage was already deferred when it was initially - encountered, and is now in the midst of executing early, so the new - deferred grouped fields set can be executed immediately. - """ - return parent_defer_usages is None or not any( - defer_usage in parent_defer_usages for defer_usage in defer_usage_set - ) - - -def build_completed_execution_group( - errors: list[GraphQLError] | None, - pending_execution_group: PendingExecutionGroup, - path: Path | None, - result: GraphQLWrappedResult[dict[str, Any]], -) -> CompletedExecutionGroup: - """Build a completed execution group.""" - return SuccessfulExecutionGroup( - pending_execution_group=pending_execution_group, - path=path.as_list() if path else [], - result=ExecutionGroupResult(result.result, errors or None), - incremental_data_records=result.increments, - ) - - -def get_deferred_fragment_records( - defer_usages: DeferUsageSet, defer_map: RefMap[DeferUsage, DeferredFragmentRecord] -) -> list[DeferredFragmentRecord]: - """Get the deferred fragment records for the given defer usages.""" - return [ - deferred_fragment_record_from_defer_usage(defer_usage, defer_map) - for defer_usage in defer_usages - ] - - -def build_stream_item_result( - result: GraphQLWrappedResult[Any], errors: list[GraphQLError] | None = None -) -> StreamItemResult: - """Build a stream item result.""" - return StreamItemResult(result.result, result.increments, errors) - - -def get_typename(value: Any) -> str | None: - """Get the ``__typename`` property of the given value.""" - if isinstance(value, Mapping): - return value.get("__typename") - # need to de-mangle the attribute assumed to be "private" in Python - for cls in value.__class__.__mro__: - __typename = getattr(value, f"_{cls.__name__}__typename", None) - if __typename: - return __typename - return None - - -def default_type_resolver( - value: Any, info: GraphQLResolveInfo, abstract_type: GraphQLAbstractType -) -> AwaitableOrValue[str | None]: - """Default type resolver function. - - If a resolve_type function is not given, then a default resolve behavior is used - which attempts two strategies: - - First, See if the provided value has a ``__typename`` field defined, if so, use that - value as name of the resolved type. - - Otherwise, test each possible type for the abstract type by calling - :meth:`~graphql.type.GraphQLObjectType.is_type_of` for the object - being coerced, returning the first type that matches. - """ - # First, look for `__typename`. - type_name = get_typename(value) - if isinstance(type_name, str): - return type_name - - # Otherwise, test each possible type. - possible_types = info.schema.get_possible_types(abstract_type) - is_awaitable = info.is_awaitable - awaitable_is_type_of_results: list[Awaitable[bool]] = [] - append_awaitable_result = awaitable_is_type_of_results.append - awaitable_types: list[GraphQLObjectType] = [] - append_awaitable_type = awaitable_types.append - - try: - for type_ in possible_types: - if type_.is_type_of: - is_type_of_result = type_.is_type_of(value, info) - - if is_awaitable(is_type_of_result): - append_awaitable_result(cast("Awaitable[bool]", is_type_of_result)) - append_awaitable_type(type_) - elif is_type_of_result: - if awaitable_is_type_of_results: - info.async_helpers.track(awaitable_is_type_of_results) - return type_.name - except Exception: - if awaitable_is_type_of_results: - # Settle the pending isTypeOf results in the background so that - # their errors can be observed before they would be orphaned. - info.async_helpers.track(awaitable_is_type_of_results) - raise - - if awaitable_is_type_of_results: - - async def get_type() -> str | None: - is_type_of_results = await info.async_helpers.gather( - awaitable_is_type_of_results - ) - for is_type_of_result, type_ in zip( - is_type_of_results, awaitable_types, strict=True - ): - if is_type_of_result: - return type_.name - return None - - return get_type() - - return None - - -def default_field_resolver(source: Any, info: GraphQLResolveInfo, **args: Any) -> Any: - """Default field resolver. - - If a resolve function is not given, then a default resolve behavior is used which - takes the property of the source object of the same name as the field and returns - it as the result, or if it's a function, returns the result of calling that function - while passing along args and context. - - For dictionaries, the field names are used as keys, for all other objects they are - used as attribute names. - """ - # Ensure source is a value for which property access is acceptable. - field_name = info.field_name - value = ( - source.get(field_name) - if isinstance(source, Mapping) - else getattr(source, field_name, None) - ) - if callable(value): - return value(info, **args) - return value - - def subscribe( schema: GraphQLSchema, document: DocumentNode, @@ -2925,7 +351,7 @@ def subscribe( custom ``root_selection_set_executor``. """ if executor_class is None: - executor_class = Executor + executor_class = ExecutorThrowingOnIncremental # If a valid executor cannot be created due to incorrect arguments, # a "Response" with only errors is returned. @@ -3005,7 +431,7 @@ def execute_subscription_event( RootSelectionSetExecutor: TypeAlias = Callable[ - ["Executor"], AwaitableOrValue[ExecutionResult] + ["Executor"], "AwaitableOrValue[ExecutionResult]" ] @@ -3179,17 +605,3 @@ def assert_event_stream(result: Any) -> AsyncIterable: raise GraphQLError(msg) return result - - -def collect_iterator_awaitables( - iterator: Iterator[Any], is_awaitable: Callable[[Any], bool] -) -> list[Awaitable[Any]]: - """Drain a synchronous iterator after an abrupt completion. - - Collects any awaitable values the iterator still holds so that their errors - can be observed before they would otherwise be left orphaned. - """ - awaitables: list[Awaitable[Any]] = [] - with suppress_exceptions: - awaitables.extend(item for item in iterator if is_awaitable(item)) - return awaitables diff --git a/src/graphql/execution/executor.py b/src/graphql/execution/executor.py new file mode 100644 index 00000000..d32e837d --- /dev/null +++ b/src/graphql/execution/executor.py @@ -0,0 +1,1969 @@ +"""GraphQL executor""" + +from __future__ import annotations + +from asyncio import ( + FIRST_COMPLETED, + ensure_future, + gather, + get_running_loop, + iscoroutine, + wait, +) +from collections.abc import ( + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Iterable, + Iterator, + Mapping, + Sequence, +) +from contextlib import suppress +from copy import copy +from typing import ( + TYPE_CHECKING, + Any, + Generic, + NamedTuple, + TypeVar, + cast, +) + +from ..error import GraphQLError, located_error +from ..language import ( + DocumentNode, + FieldNode, + FragmentDefinitionNode, + OperationDefinitionNode, + OperationType, +) +from ..pyutils import ( + AbortSignal, + AwaitableOrValue, + Path, + RefMap, + Undefined, + async_reduce, + gather_with_cancel, + inspect, + is_iterable, +) +from ..pyutils.is_awaitable import ( + is_async_iterable as default_is_async_iterable, +) +from ..pyutils.is_awaitable import ( + is_awaitable as default_is_awaitable, +) +from ..type import ( + GraphQLAbstractType, + GraphQLField, + GraphQLFieldResolver, + GraphQLLeafType, + GraphQLList, + GraphQLObjectType, + GraphQLOutputType, + GraphQLResolveInfo, + GraphQLResolveInfoHelpers, + GraphQLSchema, + GraphQLStreamDirective, + GraphQLTypeResolver, + assert_valid_schema, + is_abstract_type, + is_leaf_type, + is_list_type, + is_non_null_type, + is_object_type, +) +from ..type.directives import GraphQLDisableErrorPropagationDirective +from .aborted_graphql_execution_error import AbortedGraphQLExecutionError +from .collect_fields import ( + CollectedFields, + DeferUsage, + FieldDetails, + FieldDetailsList, + FragmentDetails, + GroupedFieldSet, + collect_fields, + collect_subfields, +) +from .get_variable_signature import get_variable_signature +from .middleware import MiddlewareManager +from .types import ExecutionResult, ExperimentalIncrementalExecutionResults +from .values import ( + VariableValues, + get_argument_values, + get_directive_values, + get_variable_values, +) + +if TYPE_CHECKING: + from asyncio import Future + from typing import TypeAlias, TypeGuard + + from ..pyutils import UndefinedType + from .get_variable_signature import GraphQLVariableSignature + +__all__ = [ + "AsyncWorkFinishedInfo", + "CollectedErrors", + "ExecutionHooks", + "Executor", + "Middleware", + "StreamUsage", + "default_field_resolver", + "default_type_resolver", +] + +T = TypeVar("T") +TContext = TypeVar("TContext") + +suppress_exceptions = suppress(Exception) + +UNEXPECTED_MULTIPLE_PAYLOADS = ( + "Executing this GraphQL operation would unexpectedly produce multiple payloads" + " (due to @defer or @stream directive)" +) + +DEFER_NOT_SUPPORTED_ON_SUBSCRIPTIONS = ( + "`@defer` directive not supported on subscription operations." + " Disable `@defer` by setting the `if` argument to `false`." +) + +Middleware: TypeAlias = tuple | list | MiddlewareManager | None + + +class StreamUsage(NamedTuple): + """Stream directive usage information""" + + label: str | None + initial_count: int + field_details_list: FieldDetailsList + + +class AsyncWorkFinishedInfo(NamedTuple): + """Information passed to the ``async_work_finished`` execution hook.""" + + executor: Executor + + +class ExecutionHooks(NamedTuple): + """Hooks for observing the execution of a GraphQL operation. + + The ``async_work_finished`` hook is run when all asynchronous work tracked + by the execution has finished. Cancelled asynchronous work may still be + running even after the result has been delivered; this hook allows + interested execution harnesses to track when this asynchronous work + completes. Errors raised by the hook are ignored. + """ + + async_work_finished: Callable[[AsyncWorkFinishedInfo], None] | None = None + + +class CollectedErrors: + """Field errors collected during execution, tracking nulled positions. + + For internal use only. + """ + + __slots__ = "_error_positions", "_errors" + + _error_positions: set[Path | None] + _errors: list[GraphQLError] + + def __init__(self) -> None: + self._error_positions = set() + self._errors = [] + + @property + def errors(self) -> list[GraphQLError]: + """Get the collected errors.""" + return self._errors + + def add(self, error: GraphQLError, path: Path | None) -> None: + """Add an error raised at the given execution position. + + Does not modify the errors list if the execution position for this + error or any of its ancestors has already been nulled via error + propagation. This check should be unnecessary for implementations + able to implement actual cancellation. + """ + if self.has_nulled_position(path): + return + self._error_positions.add(path) + self._errors.append(error) + + def has_nulled_position(self, start_path: Path | None) -> bool: + """Check whether the given position or one of its ancestors is nulled.""" + error_positions = self._error_positions + path = start_path + while path is not None: + if path in error_positions: + return True + path = path.prev + return None in error_positions + + +class Executor(Generic[TContext]): + """Executor for a validated GraphQL operation. + + Carries the data that must be available at all points during query + execution - namely, the schema of the type system that is currently + executing and the fragments defined in the query document - together + with the state of the current execution. + + This base executor implements plain execution without incremental + delivery: any ``@defer`` and ``@stream`` directives in the operation + are ignored. + """ + + schema: GraphQLSchema + # TODO: consider deprecating/removing fragment_definitions if/when fragment + # arguments are officially supported and/or the full fragment details are + # exposed within GraphQLResolveInfo. + fragment_definitions: dict[str, FragmentDefinitionNode] + fragments: dict[str, FragmentDetails] + root_value: Any + context_value: Any + operation: OperationDefinitionNode + variable_values: VariableValues + field_resolver: GraphQLFieldResolver + type_resolver: GraphQLTypeResolver + subscribe_field_resolver: GraphQLFieldResolver + enable_early_execution: bool + hide_suggestions: bool + abort_signal: AbortSignal | None + hooks: ExecutionHooks | None + async_helpers: GraphQLResolveInfoHelpers + collected_errors: CollectedErrors + pending_incremental_futures: set[Future[Any]] + background_futures: set[Future[Any]] + async_work_finished_hook_task: Future[None] | None + middleware_manager: MiddlewareManager | None + error_propagation: bool + + is_awaitable: Callable[[Any], TypeGuard[Awaitable]] = staticmethod( + default_is_awaitable # type: ignore + ) + is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] = staticmethod( + default_is_async_iterable # type: ignore + ) + + def __init__( # noqa: PLR0913 + self, + schema: GraphQLSchema, + fragment_definitions: dict[str, FragmentDefinitionNode], + fragments: dict[str, FragmentDetails], + root_value: Any, + context_value: Any, + operation: OperationDefinitionNode, + variable_values: VariableValues, + field_resolver: GraphQLFieldResolver, + type_resolver: GraphQLTypeResolver, + subscribe_field_resolver: GraphQLFieldResolver, + enable_early_execution: bool = False, + middleware_manager: MiddlewareManager | None = None, + is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None, + is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None, + hide_suggestions: bool = False, + abort_signal: AbortSignal | None = None, + hooks: ExecutionHooks | None = None, + ) -> None: + self.schema = schema + self.fragment_definitions = fragment_definitions + self.fragments = fragments + self.root_value = root_value + self.context_value = context_value + self.operation = operation + self.variable_values = variable_values + self.field_resolver = field_resolver + self.type_resolver = type_resolver + self.subscribe_field_resolver = subscribe_field_resolver + self.enable_early_execution = enable_early_execution + self.hide_suggestions = hide_suggestions + self.abort_signal = abort_signal + self.hooks = hooks + self.async_helpers = GraphQLResolveInfoHelpers( + gather=self.gather_async_work, track=self.track_async_work + ) + self.middleware_manager = middleware_manager + self.error_propagation = not any( + directive.name.value == GraphQLDisableErrorPropagationDirective.name + for directive in operation.directives or () + ) + self.is_awaitable = is_awaitable or default_is_awaitable + self.is_async_iterable = is_async_iterable or default_is_async_iterable + self.collected_errors = CollectedErrors() + # Attributes holding shared execution state: these are shared between + # this executor and all its sub-executors, since sub-executors are + # created as shallow copies that replace only the per-execution state. + self.pending_incremental_futures = set() + self.background_futures = set() + self.async_work_finished_hook_task = None + self._relevant_sub_fields: dict[tuple, CollectedFields] = {} + self._stream_usages: RefMap[FieldDetailsList, StreamUsage] = RefMap() + + @classmethod + def build( # noqa: PLR0913 + cls, + schema: GraphQLSchema, + document: DocumentNode, + root_value: Any = None, + context_value: Any = None, + raw_variable_values: dict[str, Any] | None = None, + operation_name: str | None = None, + field_resolver: GraphQLFieldResolver | None = None, + type_resolver: GraphQLTypeResolver | None = None, + subscribe_field_resolver: GraphQLFieldResolver | None = None, + max_coercion_errors: int = 50, + enable_early_execution: bool = False, + middleware: Middleware | None = None, + is_awaitable: Callable[[Any], TypeGuard[Awaitable]] | None = None, + is_async_iterable: Callable[[Any], TypeGuard[AsyncIterable]] | None = None, + hide_suggestions: bool = False, + abort_signal: AbortSignal | None = None, + hooks: ExecutionHooks | None = None, + **custom_args: Any, + ) -> list[GraphQLError] | Executor: + """Build an executor + + Constructs an Executor object from the arguments passed to execute, which + we will pass throughout the other execution methods. + + Throws a GraphQLError if a valid executor cannot be created. + + For internal use only. + """ + # If the schema used for execution is invalid, raise an error. + assert_valid_schema(schema) + + operation: OperationDefinitionNode | None = None + fragment_definitions: dict[str, FragmentDefinitionNode] = {} + fragments: dict[str, FragmentDetails] = {} + fragment_variable_signature_errors: list[GraphQLError] = [] + middleware_manager: MiddlewareManager | None = None + if middleware is not None: + if isinstance(middleware, (list, tuple)): + middleware_manager = MiddlewareManager(*middleware) + elif isinstance(middleware, MiddlewareManager): + middleware_manager = middleware + else: + msg = ( + "Middleware must be passed as a list or tuple of functions" + " or objects, or as a single MiddlewareManager object." + f" Got {inspect(middleware)} instead." + ) + raise TypeError(msg) + + for definition in document.definitions: + if isinstance(definition, OperationDefinitionNode): + if operation_name is None: + if operation: + return [ + GraphQLError( + "Must provide operation name" + " if query contains multiple operations." + ) + ] + operation = definition + elif definition.name and definition.name.value == operation_name: + operation = definition + elif isinstance(definition, FragmentDefinitionNode): + fragment_definitions[definition.name.value] = definition + variable_signatures: dict[str, GraphQLVariableSignature] | None = None + if definition.variable_definitions: + variable_signatures = {} + for var_def in definition.variable_definitions: + signature = get_variable_signature(schema, var_def) + if isinstance(signature, GraphQLError): + fragment_variable_signature_errors.append(signature) + continue + variable_signatures[signature.name] = signature + fragments[definition.name.value] = FragmentDetails( + definition, variable_signatures + ) + + if not operation: + if operation_name is not None: + return [GraphQLError(f"Unknown operation named '{operation_name}'.")] + return [GraphQLError("Must provide an operation.")] + + if fragment_variable_signature_errors: + return fragment_variable_signature_errors + + variable_values = get_variable_values( + schema, + operation.variable_definitions or (), + raw_variable_values or {}, + max_errors=max_coercion_errors, + hide_suggestions=hide_suggestions, + ) + + if isinstance(variable_values, list): + return variable_values # errors + + return cls( + schema, + fragment_definitions, + fragments, + root_value, + context_value, + operation, + variable_values, + field_resolver or default_field_resolver, + type_resolver or default_type_resolver, + subscribe_field_resolver or default_field_resolver, + enable_early_execution, + middleware_manager, + is_awaitable, + is_async_iterable, + hide_suggestions=hide_suggestions, + abort_signal=abort_signal, + hooks=hooks, + **custom_args, + ) + + def build_per_event_executor(self, payload: Any) -> Executor: + """Create a copy of the executor for usage with subscribe events.""" + executor = copy(self) + executor.root_value = payload + executor.collected_errors = CollectedErrors() + return executor + + def execute_operation( + self, + serially: bool | None = None, + ) -> AwaitableOrValue[ExecutionResult | ExperimentalIncrementalExecutionResults]: + """Execute an operation. + + Implements the "Executing operations" section of the spec. + + Return a possible coroutine object that will eventually yield the data described + by the "Response" section of the GraphQL specification. + + If errors are encountered while executing a GraphQL field, only that field and + its descendants will be omitted, and sibling fields will still be executed. An + execution which encounters errors will still result in a coroutine object that + can be executed without errors. + + Errors from sub-fields of a NonNull type may propagate to the top level, + at which point we still log the error and null the parent field, which + in this case is the entire response. + + If the operation is aborted, the whole operation is rejected with an + aborted execution error rather than resolving to a partial response with + located errors; the partial result that the unwinding execution can still + produce is exposed on that error. + """ + abort_signal = self.abort_signal + if abort_signal is not None and abort_signal.aborted: + self.run_async_work_finished_hook() + raise self.abort_error() + try: + operation = self.operation + schema = self.schema + operation_type = operation.operation + root_type = schema.get_root_type(operation_type) + if root_type is None: + msg = ( + "Schema is not configured to execute" + f" {operation_type.value} operation." + ) + raise GraphQLError(msg, operation) # noqa: TRY301 + root_value = self.root_value + + collected_fields = collect_fields( + schema, + self.fragments, + self.variable_values, + root_type, + operation, + self.hide_suggestions, + ) + + grouped_field_set, new_defer_usages, _forbidden = collected_fields + + result = self.execute_collected_root_fields( + root_type, + root_value, + grouped_field_set, + operation_type == OperationType.MUTATION + if serially is None + else serially, + new_defer_usages, + ) + + if self.is_awaitable(result): + + async def await_result() -> ( + ExecutionResult | ExperimentalIncrementalExecutionResults + ): + try: + data = await result + except GraphQLError as error: + self.collected_errors.add(error, None) + return self.build_response(None) + except Exception: + # cancel incremental work started early and close the + # stream sources before re-raising, e.g. the abort reason + await self.cancel_incremental_work() + self.run_async_work_finished_hook() + raise + return self.build_response(data) + + if abort_signal is None: + return await_result() + return self.with_aborted_execution_error(await_result()) + + data = cast("dict[str, Any]", result) + + except GraphQLError as error: + self.collected_errors.add(error, None) + return self.finish(self.build_response(None)) + + return self.finish(self.build_response(data)) + + def execute_collected_root_fields( + self, + root_type: GraphQLObjectType, + root_value: Any, + grouped_field_set: GroupedFieldSet, + serially: bool, + new_defer_usages: Sequence[DeferUsage], + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the collected root fields, ignoring incremental delivery.""" + return self.execute_root_grouped_field_set( + root_type, + root_value, + grouped_field_set, + serially, + None, + ) + + def execute_root_grouped_field_set( + self, + root_type: GraphQLObjectType, + root_value: Any, + grouped_field_set: GroupedFieldSet, + serially: bool, + position_context: TContext | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the root grouped field set.""" + return (self.execute_fields_serially if serially else self.execute_fields)( + root_type, + root_value, + None, + grouped_field_set, + position_context, + ) + + def build_response( + self, data: dict[str, Any] | None + ) -> ExecutionResult | ExperimentalIncrementalExecutionResults: + """Build the response for the given completed data. + + Given completed execution data, build the ``(data, errors)`` response + defined by the "Response" section of the GraphQL specification. + """ + self.run_async_work_finished_hook() + errors = self.collected_errors.errors + return ExecutionResult(data, errors or None) + + def execute_fields_serially( + self, + parent_type: GraphQLObjectType, + source_value: Any, + path: Path | None, + grouped_field_set: GroupedFieldSet, + position_context: TContext | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the given fields serially. + + Implements the "Executing selection sets" section of the spec + for fields that must be executed serially. + """ + is_awaitable = self.is_awaitable + abort_signal = self.abort_signal + + def reducer( + results: dict[str, Any], + field_item: tuple[str, FieldDetailsList], + ) -> AwaitableOrValue[dict[str, Any]]: + response_name, field_details_list = field_item + field_path = Path(path, response_name, parent_type.name) + if abort_signal is not None and abort_signal.aborted: + # Reject the whole operation rather than serially completing the + # remaining fields with located abort errors. + raise self.abort_error() + result = self.execute_field( + parent_type, + source_value, + field_details_list, + field_path, + position_context, + ) + if result is Undefined: + return results + if is_awaitable(result): + + async def set_result() -> dict[str, Any]: + results[response_name] = await result + return results + + return set_result() + + results[response_name] = result + return results + + return async_reduce(reducer, grouped_field_set.items(), {}) + + def execute_fields( + self, + parent_type: GraphQLObjectType, + source_value: Any, + path: Path | None, + grouped_field_set: GroupedFieldSet, + position_context: TContext | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the given fields concurrently. + + Implements the "Executing selection sets" section of the spec + for fields that may be executed in parallel. + """ + results: dict[str, Any] = {} + is_awaitable = self.is_awaitable + awaitable_fields: list[str] = [] + append_awaitable = awaitable_fields.append + try: + for response_name, field_details_list in grouped_field_set.items(): + field_path = Path(path, response_name, parent_type.name) + result = self.execute_field( + parent_type, + source_value, + field_details_list, + field_path, + position_context, + ) + if result is not Undefined: + results[response_name] = result + if is_awaitable(result): + append_awaitable(response_name) + except Exception: + if awaitable_fields: + # Ensure that awaitables created by other fields are settled, + # as they may also fail. + self.settle_in_background( + [results[field] for field in awaitable_fields] + ) + raise + + # If there are no coroutines, we can just return the object. + if not awaitable_fields: + return results + + # Otherwise, results is a map from field name to the result of resolving that + # field, which is possibly a coroutine object. Return a coroutine object that + # will yield this same map, but with any coroutines awaited in parallel and + # replaced with the values they yielded. + async def get_results() -> dict[str, Any]: + if len(awaitable_fields) == 1: + # If there is only one field, avoid the overhead of parallelization. + field = awaitable_fields[0] + results[field] = await results[field] + else: + awaited_results = await gather_with_cancel( + *(results[field] for field in awaitable_fields) + ) + results.update(zip(awaitable_fields, awaited_results, strict=True)) + + return results + + return get_results() + + def execute_field( + self, + parent_type: GraphQLObjectType, + source: Any, + field_details_list: FieldDetailsList, + path: Path, + position_context: TContext | None, + ) -> AwaitableOrValue[Any] | UndefinedType: + """Resolve the field on the given source object. + + Implements the "Executing fields" section of the spec. + + In particular, this method figures out the value that the field returns by + calling its resolve function, then calls complete_value to await coroutine + objects, coercing scalars, or execute the sub-selection-set for objects. + """ + first_field_details = field_details_list[0] + first_field_node = first_field_details.node + field_name = first_field_node.name.value + field_def = self.schema.get_field(parent_type, field_name) + if not field_def: + return Undefined + + return_type = field_def.type + resolve_fn = field_def.resolve or self.field_resolver + + if self.middleware_manager: + resolve_fn = self.middleware_manager.get_field_resolver(resolve_fn) + + info = self.build_resolve_info( + field_def, to_nodes(field_details_list), parent_type, path + ) + + # Get the resolve function, regardless of if its result is normal or abrupt + # (error). + try: + # Build a dictionary of arguments from the field.arguments AST, using the + # variables scope to fulfill any variable references. + args = get_argument_values( + field_def, + first_field_node, + self.variable_values, + first_field_details.fragment_variable_values, + self.hide_suggestions, + ) + + # Note that contrary to the JavaScript implementation, we pass the context + # value as part of the resolve info. + result = resolve_fn(source, info, **args) + + if self.is_awaitable(result): + return self.complete_awaitable_value( + return_type, + field_details_list, + info, + path, + result, + position_context, + ) + + completed = self.complete_value( + return_type, + field_details_list, + info, + path, + result, + position_context, + ) + if self.is_awaitable(completed): + + async def await_completed() -> Any: + try: + return await completed + except Exception as raw_error: + self.handle_field_error( + raw_error, + return_type, + field_details_list, + path, + ) + return None + + return await_completed() + + except Exception as raw_error: + self.handle_field_error( + raw_error, + return_type, + field_details_list, + path, + ) + return None + + return completed + + def build_resolve_info( + self, + field_def: GraphQLField, + field_nodes: list[FieldNode], + parent_type: GraphQLObjectType, + path: Path, + ) -> GraphQLResolveInfo: + """Build the GraphQLResolveInfo object. + + For internal use only. + """ + # The resolve function's first argument is a collection of information about + # the current execution state. + return GraphQLResolveInfo( + field_nodes[0].name.value, + field_nodes, + field_def.type, + parent_type, + path, + self.schema, + self.fragment_definitions, + self.root_value, + self.operation, + self.variable_values, + self.context_value, + self.is_awaitable, + self.abort_signal, + self.async_helpers, + ) + + def handle_field_error( + self, + raw_error: Exception, + return_type: GraphQLOutputType, + field_details_list: FieldDetailsList, + path: Path, + ) -> None: + """Handle error properly according to the field type.""" + error = located_error(raw_error, to_nodes(field_details_list), path.as_list()) + + # If the field type is non-nullable, then it is resolved without any protection + # from errors, however it still properly locates the error. + if self.error_propagation and is_non_null_type(return_type): + raise error + + # Otherwise, error protection is applied, logging the error and resolving a + # null value for this field if one is encountered. + self.collected_errors.add(error, path) + + def complete_value( + self, + return_type: GraphQLOutputType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + result: Any, + position_context: TContext | None, + ) -> AwaitableOrValue[Any]: + """Complete a value. + + Implements the instructions for completeValue as defined in the + "Value completion" section of the spec. + + If the field type is Non-Null, then this recursively completes the value + for the inner type. It throws a field error if that completion returns null, + as per the "Nullability" section of the spec. + + If the field type is a List, then this recursively completes the value + for the inner type on each item in the list. + + If the field type is a Scalar or Enum, ensures the completed value is a legal + value of the type by calling the ``coerce_output_value`` method of GraphQL type + definition. + + If the field is an abstract type, determine the runtime type of the value and + then complete based on that type. + + Otherwise, the field type expects a sub-selection set, and will complete the + value by evaluating all sub-selections. + """ + # If result is an Exception, throw a located error. + if isinstance(result, Exception): + raise result + + # If field type is NonNull, complete for inner type, and throw field error if + # result is null. + if is_non_null_type(return_type): + completed = self.complete_value( + return_type.of_type, + field_details_list, + info, + path, + result, + position_context, + ) + if completed is None: + msg = ( + "Cannot return null for non-nullable field" + f" {info.parent_type}.{info.field_name}." + ) + raise TypeError(msg) + return completed + + # If result value is null or undefined then return null. + if result is None or result is Undefined: + return None + + # If field type is List, complete each item in the list with inner type + if is_list_type(return_type): + return self.complete_list_value( + return_type, + field_details_list, + info, + path, + result, + position_context, + ) + + # If field type is a leaf type, Scalar or Enum, coerce to a valid value, + # returning null if coercion is not possible. + if is_leaf_type(return_type): + return self.complete_leaf_value(return_type, result) + + # If field type is an abstract type, Interface or Union, determine the runtime + # Object type and complete for that type. + if is_abstract_type(return_type): + return self.complete_abstract_value( + return_type, + field_details_list, + info, + path, + result, + position_context, + ) + + # If field type is Object, execute and complete all sub-selections. + if is_object_type(return_type): + return self.complete_object_value( + return_type, + field_details_list, + info, + path, + result, + position_context, + ) + + # Not reachable. All possible output types have been considered. + msg = ( + "Cannot complete value of unexpected output type:" + f" '{inspect(return_type)}'." + ) # pragma: no cover + raise TypeError(msg) # pragma: no cover + + async def with_abort_signal(self, awaitable: Awaitable[T]) -> T: + """Await a value, but cancel immediately if the abort signal is triggered. + + This wraps awaitables returned by resolvers (and awaitable list items) so + that a triggered abort signal interrupts execution *immediately* instead of + only at the next field boundary. Without this, a hanging asynchronous + resolver would prevent the operation from ever being cancelled. + + If the abort signal fires before the awaitable settles, the underlying + awaitable is cancelled and the abort reason is raised (an exception reason + is raised as is, any other value is reported as an unexpected error value). + """ + abort_signal = self.abort_signal + if abort_signal is None: + return await awaitable + task = ensure_future(awaitable) + if not abort_signal.aborted: + abort = ensure_future(abort_signal.wait()) + try: + await wait({task, abort}, return_when=FIRST_COMPLETED) + finally: + if not abort.done(): + abort.cancel() + if not abort_signal.aborted: + return task.result() + # The abort signal fired (possibly in the same tick the task settled); + # discard any task result and reject with the abort reason. + task.cancel() + with suppress(BaseException): + await task + raise self.abort_error() + + def abort_error(self) -> Exception: + """Return the exception to raise when execution has been aborted. + + An abort reason that is itself an exception is raised as is; any other value + is reported as an unexpected error value. + """ + reason = self.abort_signal.reason # type: ignore[union-attr] + if isinstance(reason, Exception): + return reason + msg = f"Unexpected error value: {inspect(reason)}" + return TypeError(msg) + + async def with_aborted_execution_error(self, awaitable: Awaitable[T]) -> T: + """Await a result, but raise an aborted execution error when aborted. + + Unlike :meth:`with_abort_signal`, the awaited work is not cancelled as a + whole when the abort signal is triggered: only its in-flight resolvers are + cancelled individually, so that it still settles into a partial result, + which is exposed via the raised aborted execution error. + """ + abort_signal = self.abort_signal + task = ensure_future(awaitable) + if not abort_signal.aborted: # type: ignore[union-attr] + abort = ensure_future(abort_signal.wait()) # type: ignore[union-attr] + try: + await wait({task, abort}, return_when=FIRST_COMPLETED) + finally: + if not abort.done(): + abort.cancel() + if not abort_signal.aborted: # type: ignore[union-attr] + return task.result() + # The abort signal fired (possibly in the same tick the task settled); + # let the partial result settle in the background and expose it on the + # aborted execution error that the operation is rejected with. + self.settle_in_background([task]) + raise self.create_aborted_execution_error(task) + + def finish(self, result: T) -> T: + """Check that execution has not been aborted before returning its result. + + If the operation has been aborted during otherwise synchronous execution, + raise an aborted execution error exposing the already completed result. + """ + abort_signal = self.abort_signal + if abort_signal is not None and abort_signal.aborted: + raise self.create_aborted_execution_error(result) + return result + + def create_aborted_execution_error( + self, result: AwaitableOrValue[Any] + ) -> AbortedGraphQLExecutionError: + """Create an aborted execution error exposing the given result.""" + reason = self.abort_signal.reason # type: ignore[union-attr] + return AbortedGraphQLExecutionError(reason, result) + + def abort(self, reason: BaseException | None = None) -> AwaitableOrValue[None]: + """Abort the incremental work produced by this executor. + + The base executor produces no incremental work, so this is a no-op. + """ + + def track_incremental_future(self, future: Future[Any]) -> None: + """Register a pending future belonging to incremental work. + + The registered futures can be cancelled via + :meth:`cancel_incremental_work` when the incremental execution is + stopped before they have settled. + """ + futures = self.pending_incremental_futures + futures.add(future) + future.add_done_callback(futures.discard) + + async def cancel_incremental_work( + self, reason: BaseException | None = None + ) -> None: + """Cancel all pending incremental work and close the stream sources. + + Aborts the incremental work produced by this executor, which cancels + the still pending incremental execution tasks and triggers the early + return of the stream sources, and cancels any remaining pending + incremental futures, waiting until all cancellation has settled. + """ + abort_result = self.abort(reason) + if default_is_awaitable(abort_result): + await abort_result + futures = self.pending_incremental_futures + if futures: + pending = list(futures) + for future in pending: + future.cancel() + await gather(*pending, return_exceptions=True) + + def settle_in_background(self, awaitables: list[Awaitable[Any]]) -> None: + """Settle the given pending awaitables in the background. + + A bubbling synchronous error must not wait for pending sibling awaitables, + but they must still be settled so that their errors can be observed before + they would be orphaned (the Python analog of silencing JS unhandled + rejections). Without a running event loop the awaitables can never run; + pending coroutines are then closed instead to dispose of them. + """ + try: + get_running_loop() + except RuntimeError: # no running event loop + for awaitable in awaitables: + if iscoroutine(awaitable): # pragma: no branch + awaitable.close() + return + future = gather(*awaitables, return_exceptions=True) + background_futures = self.background_futures + background_futures.add(future) + future.add_done_callback(background_futures.discard) + + def track_async_work(self, values: Sequence[Any]) -> None: + """Track possibly awaitable values as pending asynchronous work. + + Awaitables among the given values are settled in the background, so that + they are still settled and their errors observed when they would otherwise + be abandoned. Non-awaitable values are ignored. + """ + is_awaitable = self.is_awaitable + awaitables: list[Awaitable[Any]] = [ + value for value in values if is_awaitable(value) + ] + if awaitables: + self.settle_in_background(awaitables) + + def gather_async_work( + self, values: Sequence[Awaitable[Any]] + ) -> Awaitable[list[Any]]: + """Concurrently await the given values as one unit of asynchronous work. + + This allows resolvers to await multiple concurrent operations together. + When one of the values fails, the others are cancelled and settled before + the error is propagated, so that no asynchronous work is orphaned. + """ + return gather_with_cancel(*values) + + def run_async_work_finished_hook(self) -> None: + """Run the hook signaling that all asynchronous work has finished. + + If an ``async_work_finished`` execution hook is provided, run it as soon as + all tracked pending asynchronous work has been settled - synchronously when + there is no pending asynchronous work, which allows synchronous execution + paths to remain synchronous. Errors raised by the hook are ignored. + """ + hooks = self.hooks + hook = hooks.async_work_finished if hooks is not None else None + if hook is None: + return + info = AsyncWorkFinishedInfo(self) + background_futures = self.background_futures + if not background_futures: + with suppress_exceptions: + hook(info) + return + + async def wait_and_run_hook() -> None: + while background_futures: + await wait(list(background_futures)) + with suppress_exceptions: + hook(info) + + # keep a reference to the task so that it is not garbage collected + self.async_work_finished_hook_task = ensure_future(wait_and_run_hook()) + + def cancellable_iterable(self, iterable: AsyncIterable[T]) -> AsyncIterable[T]: + """Wrap an async iterable so pending iteration is cancelled on abort. + + When the abort signal is triggered, any pending ``__anext__`` call returns + immediately by raising the abort reason. This mirrors the JavaScript + ``cancellableIterable``; GraphQL-core needs no ``AbortSignalListener`` + class since :meth:`with_abort_signal` already provides the cancellation + mechanism. + """ + if self.abort_signal is None: + return iterable + with_abort_signal = self.with_abort_signal + iterator = iterable.__aiter__() + + class CancellableAsyncIterator: + def __aiter__(self) -> AsyncIterator[T]: + return self + + def __anext__(self) -> Awaitable[T]: + return with_abort_signal(iterator.__anext__()) + + async def aclose(self) -> None: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + + return CancellableAsyncIterator() + + async def complete_awaitable_value( + self, + return_type: GraphQLOutputType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + result: Any, + position_context: TContext | None, + ) -> Any: + """Complete an awaitable value.""" + try: + resolved = await self.with_abort_signal(result) + completed = self.complete_value( + return_type, + field_details_list, + info, + path, + resolved, + position_context, + ) + if self.is_awaitable(completed): + completed = await completed + except Exception as raw_error: + self.handle_field_error(raw_error, return_type, field_details_list, path) + completed = None + return completed + + def get_stream_usage( + self, field_details_list: FieldDetailsList, path: Path + ) -> StreamUsage | None: + """Get stream usage. + + Returns an object containing info for streaming if a field should be + streamed based on the experimental flag, stream directive present and + not disabled by the "if" argument. + """ + # do not stream inner lists of multidimensional lists + if isinstance(path.key, int): + return None + + stream_usage = self._stream_usages.get(field_details_list) + if stream_usage is not None: + return stream_usage # pragma: no cover + + # validation only allows equivalent streams on multiple fields, so it is + # safe to only check the first field_node for the stream directive + stream = get_directive_values( + GraphQLStreamDirective, field_details_list[0].node, self.variable_values + ) + + if not stream or stream.get("if") is False: + return None + + initial_count = stream.get("initialCount") + if initial_count is None or initial_count < 0: + msg = "initialCount must be a positive integer" + raise ValueError(msg) + + if self.operation.operation == OperationType.SUBSCRIPTION: + msg = ( + "`@stream` directive not supported on subscription operations." + " Disable `@stream` by setting the `if` argument to `false`." + ) + raise TypeError(msg) + + streamed_field_details_list: FieldDetailsList = [ + FieldDetails(field_details.node, None) + for field_details in field_details_list + ] + + stream_usage = StreamUsage( + stream.get("label"), stream["initialCount"], streamed_field_details_list + ) + + self._stream_usages[field_details_list] = stream_usage + + return stream_usage + + def handle_stream( + self, + index: int, + path: Path, + iterator: Iterator[Any] | AsyncIterator[Any], + is_async: bool, + stream_usage: StreamUsage, + info: GraphQLResolveInfo, + item_type: GraphQLOutputType, + ) -> bool: + """Handle streaming of the remaining list items, if supported. + + The base executor does not support streaming, so this is a no-op + returning False, which means that the list is completed normally. + """ + return False + + async def complete_async_iterator_value( + self, + item_type: GraphQLOutputType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + async_iterator: AsyncIterator[Any], + position_context: TContext | None, + ) -> list[Any]: + """Complete an async iterator. + + Complete an async iterator value by completing the result and calling + recursively until all the results are completed. + """ + is_awaitable = self.is_awaitable + complete_list_item_value = self.complete_list_item_value + complete_awaitable_list_item_value = self.complete_awaitable_list_item_value + completed_results: list[Any] = [] + append_completed = completed_results.append + awaitable_indices: list[int] = [] + append_awaitable = awaitable_indices.append + stream_usage = self.get_stream_usage(field_details_list, path) + try: + early_return = async_iterator.aclose # type: ignore[attr-defined] + except AttributeError: + early_return = None + index = 0 + try: + while True: + if ( + stream_usage + and index == stream_usage.initial_count + and self.handle_stream( + index, + path, + async_iterator, + True, + stream_usage, + info, + item_type, + ) + ): + break + + item_path = path.add_key(index, None) + try: + item = await anext(async_iterator) + except StopAsyncIteration: + break + except Exception as raw_error: + raise located_error( + raw_error, to_nodes(field_details_list), path.as_list() + ) from raw_error + + if is_awaitable(item): + append_completed( + complete_awaitable_list_item_value( + item, + item_type, + field_details_list, + info, + item_path, + position_context, + ) + ) + append_awaitable(index) + + elif complete_list_item_value( + item, + completed_results, + item_type, + field_details_list, + info, + item_path, + position_context, + ): + append_awaitable(index) + + index += 1 + except Exception: + if early_return is not None: # pragma: no branch + with suppress_exceptions: + await early_return() + if awaitable_indices: + # Settle any awaitable items already collected in the background, + # so that the current error is not delayed. + self.settle_in_background( + [completed_results[index] for index in awaitable_indices] + ) + raise + + if not awaitable_indices: + return completed_results + + if len(awaitable_indices) == 1: + # If there is only one index, avoid the overhead of parallelization. + index = awaitable_indices[0] + completed_results[index] = await completed_results[index] + else: + awaited_results = await gather_with_cancel( + *(completed_results[index] for index in awaitable_indices) + ) + for index, sub_result in zip( + awaitable_indices, awaited_results, strict=True + ): + completed_results[index] = sub_result + return completed_results + + def complete_list_value( + self, + return_type: GraphQLList[GraphQLOutputType], + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + result: AsyncIterable[Any] | Iterable[Any], + position_context: TContext | None, + ) -> AwaitableOrValue[list[Any]]: + """Complete a list value. + + Complete a list value by completing each item in the list with the inner type. + """ + item_type = return_type.of_type + + if self.is_async_iterable(result): + async_iterator = self.cancellable_iterable(result).__aiter__() + + return self.complete_async_iterator_value( + item_type, + field_details_list, + info, + path, + async_iterator, + position_context, + ) + + if not is_iterable(result): + msg = ( + "Expected Iterable, but did not find one for field" + f" '{info.parent_type}.{info.field_name}'." + ) + raise GraphQLError(msg) + + return self.complete_iterable_value( + item_type, + field_details_list, + info, + path, + result, + position_context, + ) + + def complete_iterable_value( + self, + item_type: GraphQLOutputType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + items: Iterable[Any], + position_context: TContext | None, + ) -> AwaitableOrValue[list[Any]]: + """Complete an iterable value.""" + # This is specified as a simple map, however we're optimizing the path + # where the list contains no awaitable routine objects by avoiding creating + # another awaitable object. + is_awaitable = self.is_awaitable + complete_list_item_value = self.complete_list_item_value + complete_awaitable_list_item_value = self.complete_awaitable_list_item_value + completed_results: list[Any] = [] + append_completed = completed_results.append + awaitable_indices: list[int] = [] + append_awaitable = awaitable_indices.append + stream_usage = self.get_stream_usage(field_details_list, path) + iterator = iter(items) + index = 0 + try: + while True: + if ( + stream_usage + and index == stream_usage.initial_count + and self.handle_stream( + index, + path, + iterator, + False, + stream_usage, + info, + item_type, + ) + ): + break + + try: + item = next(iterator) + except StopIteration: + break + + # No need to modify the info object containing the path, + # since from here on it is not ever accessed by resolver functions. + item_path = path.add_key(index, None) + + if is_awaitable(item): + append_completed( + complete_awaitable_list_item_value( + item, + item_type, + field_details_list, + info, + item_path, + position_context, + ) + ) + append_awaitable(index) + + elif complete_list_item_value( + item, + completed_results, + item_type, + field_details_list, + info, + item_path, + position_context, + ): + append_awaitable(index) + + index += 1 + except Exception: + # Do not close the iterator. Instead, drain it so that any awaitable + # items it still holds can be settled in the background before they + # would be orphaned. + maybe_awaitables = [completed_results[index] for index in awaitable_indices] + maybe_awaitables.extend(collect_iterator_awaitables(iterator, is_awaitable)) + if maybe_awaitables: + self.settle_in_background(maybe_awaitables) + raise + + if not awaitable_indices: + return completed_results + + async def get_completed_results() -> list[Any]: + if len(awaitable_indices) == 1: + # If there is only one index, avoid the overhead of parallelization. + index = awaitable_indices[0] + completed_results[index] = await completed_results[index] + else: + awaited_results = await gather_with_cancel( + *(completed_results[index] for index in awaitable_indices) + ) + for index, sub_result in zip( + awaitable_indices, awaited_results, strict=True + ): + completed_results[index] = sub_result + return completed_results + + return get_completed_results() + + def complete_list_item_value( + self, + item: Any, + complete_results: list[Any], + item_type: GraphQLOutputType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + item_path: Path, + position_context: TContext | None, + ) -> bool: + """Complete a list item value by adding it to the completed results. + + Returns True if the value is awaitable. + """ + is_awaitable = self.is_awaitable + + try: + completed_item = self.complete_value( + item_type, + field_details_list, + info, + item_path, + item, + position_context, + ) + + if is_awaitable(completed_item): + + async def await_completed() -> Any: + try: + return await completed_item + except Exception as raw_error: + self.handle_field_error( + raw_error, + item_type, + field_details_list, + item_path, + ) + return None + + complete_results.append(await_completed()) + return True + + complete_results.append(completed_item) + + except Exception as raw_error: + self.handle_field_error( + raw_error, + item_type, + field_details_list, + item_path, + ) + complete_results.append(None) + + return False + + async def complete_awaitable_list_item_value( + self, + item: Any, + item_type: GraphQLOutputType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + item_path: Path, + position_context: TContext | None, + ) -> Any: + """Complete an awaitable list item value.""" + try: + resolved = await self.with_abort_signal(item) + completed = self.complete_value( + item_type, + field_details_list, + info, + item_path, + resolved, + position_context, + ) + if self.is_awaitable(completed): + completed = await completed + except Exception as raw_error: + self.handle_field_error( + raw_error, + item_type, + field_details_list, + item_path, + ) + return None + return completed + + @staticmethod + def complete_leaf_value(return_type: GraphQLLeafType, result: Any) -> Any: + """Complete a leaf value. + + Complete a Scalar or Enum by coercing to a valid value, returning null if + coercion is not possible. + """ + coerced = return_type.coerce_output_value(result) + if coerced is Undefined or coerced is None: + msg = ( + f"Expected `{inspect(return_type)}.coerce_output_value(" + f"{inspect(result)})` to return non-nullable value, returned:" + f" {inspect(coerced)}" + ) + raise TypeError(msg) + return coerced + + def complete_abstract_value( + self, + return_type: GraphQLAbstractType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + result: Any, + position_context: TContext | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Complete an abstract value. + + Complete a value of an abstract type by determining the runtime object type of + that value, then complete the value for that type. + """ + resolve_type_fn = return_type.resolve_type or self.type_resolver + runtime_type = resolve_type_fn(result, info, return_type) + + if self.is_awaitable(runtime_type): + + async def await_complete_object_value() -> Any: + value = self.complete_object_value( + self.ensure_valid_runtime_type( + await self.with_abort_signal(runtime_type), # type: ignore + return_type, + field_details_list, + info, + result, + ), + field_details_list, + info, + path, + result, + position_context, + ) + if self.is_awaitable(value): + return await value + return value # pragma: no cover + + return await_complete_object_value() + runtime_type = cast("str | None", runtime_type) + + return self.complete_object_value( + self.ensure_valid_runtime_type( + runtime_type, return_type, field_details_list, info, result + ), + field_details_list, + info, + path, + result, + position_context, + ) + + def ensure_valid_runtime_type( + self, + runtime_type_name: Any, + return_type: GraphQLAbstractType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + result: Any, + ) -> GraphQLObjectType: + """Ensure that the given type is valid at runtime.""" + if runtime_type_name is None: + msg = ( + f"Abstract type '{return_type}' must resolve" + " to an Object type at runtime" + f" for field '{info.parent_type}.{info.field_name}'." + f" Either the '{return_type}' type should provide" + " a 'resolve_type' function or each possible type should provide" + " an 'is_type_of' function." + ) + raise GraphQLError(msg, to_nodes(field_details_list)) + + if not isinstance(runtime_type_name, str): + msg = ( + f"Abstract type '{return_type}' must resolve" + " to an Object type at runtime" + f" for field '{info.parent_type}.{info.field_name}' with value" + f" {inspect(result)}, received '{inspect(runtime_type_name)}'," + " which is not a valid Object type name." + ) + raise GraphQLError(msg, to_nodes(field_details_list)) + + runtime_type = self.schema.get_type(runtime_type_name) + + if runtime_type is None: + msg = ( + f"Abstract type '{return_type}' was resolved to a type" + f" '{runtime_type_name}' that does not exist inside the schema." + ) + raise GraphQLError(msg, to_nodes(field_details_list)) + + if not is_object_type(runtime_type): + msg = ( + f"Abstract type '{return_type}' was resolved" + f" to a non-object type '{runtime_type_name}'." + ) + raise GraphQLError(msg, to_nodes(field_details_list)) + + if not self.schema.is_sub_type(return_type, runtime_type): + msg = ( + f"Runtime Object type '{runtime_type}' is not a possible" + f" type for '{return_type}'." + ) + raise GraphQLError(msg, to_nodes(field_details_list)) + + return runtime_type + + def complete_object_value( + self, + return_type: GraphQLObjectType, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + result: Any, + position_context: TContext | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Complete an Object value by executing all sub-selections.""" + # If there is an `is_type_of()` predicate function, call it with the current + # result. If `is_type_of()` returns False, then raise an error rather than + # continuing execution. + if return_type.is_type_of: + is_type_of = return_type.is_type_of(result, info) + + if self.is_awaitable(is_type_of): + + async def execute_subfields_async() -> dict[str, Any]: + if not await self.with_abort_signal(is_type_of): + raise invalid_return_type_error( + return_type, result, field_details_list + ) + sub_fields = self.collect_and_execute_subfields( + return_type, + field_details_list, + path, + result, + position_context, + ) + if self.is_awaitable(sub_fields): + return await sub_fields + return cast("dict[str, Any]", sub_fields) # pragma: no cover + + return execute_subfields_async() + + if not is_type_of: + raise invalid_return_type_error(return_type, result, field_details_list) + + return self.collect_and_execute_subfields( + return_type, + field_details_list, + path, + result, + position_context, + ) + + def collect_and_execute_subfields( + self, + return_type: GraphQLObjectType, + field_details_list: FieldDetailsList, + path: Path, + result: Any, + position_context: TContext | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Collect sub-fields to execute to complete this value.""" + collected_subfields = self.collect_subfields(return_type, field_details_list) + grouped_field_set, new_defer_usages, _forbidden = collected_subfields + + return self.execute_collected_subfields( + return_type, + result, + path, + grouped_field_set, + new_defer_usages, + position_context, + ) + + def execute_collected_subfields( + self, + parent_type: GraphQLObjectType, + source_value: Any, + path: Path, + grouped_field_set: GroupedFieldSet, + new_defer_usages: Sequence[DeferUsage], + position_context: TContext | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the collected subfields, ignoring incremental delivery.""" + return self.execute_fields( + parent_type, + source_value, + path, + grouped_field_set, + None, + ) + + def collect_subfields( + self, return_type: GraphQLObjectType, field_details_list: FieldDetailsList + ) -> CollectedFields: + """Collect subfields. + + A memoized function collecting relevant subfields regarding the return type. + Memoizing ensures the subfields are not repeatedly calculated, which saves + overhead when resolving lists of values. + """ + relevant_sub_fields = self._relevant_sub_fields + # We cannot use the field_details_list itself as key for the cache, since it + # is not hashable as a list. We also do not want to use the field_details_list + # itself (converted to a tuple) as keys, since hashing them is slow. + # Therefore, we use the ids of the field_details_list items as keys. Note that + # we do not use the id of the list, since we want to hit the cache for all + # lists of the same nodes, not only for the same list of nodes. Also, the list + # id may even be reused, in which case we would get wrong results from cache. + key = ( + (return_type, id(field_details_list[0])) + if len(field_details_list) == 1 # optimize most frequent case + else (return_type, *map(id, field_details_list)) + ) + collected_fields: CollectedFields | None = relevant_sub_fields.get(key) + if collected_fields is None: + collected_fields = collect_subfields( + self.schema, + self.fragments, + self.variable_values, + self.operation, + return_type, + field_details_list, + self.hide_suggestions, + ) + relevant_sub_fields[key] = collected_fields + return collected_fields + + +def to_nodes(field_details_list: FieldDetailsList) -> list[FieldNode]: + """Convert a field group to a list of field nodes.""" + return [field_details.node for field_details in field_details_list] + + +def invalid_return_type_error( + return_type: GraphQLObjectType, result: Any, field_details_list: FieldDetailsList +) -> GraphQLError: + """Create a GraphQLError for an invalid return type.""" + return GraphQLError( + f"Expected value of type '{return_type}' but got: {inspect(result)}.", + to_nodes(field_details_list), + ) + + +def collect_iterator_awaitables( + iterator: Iterator[Any], is_awaitable: Callable[[Any], bool] +) -> list[Awaitable[Any]]: + """Drain a synchronous iterator after an abrupt completion. + + Collects any awaitable values the iterator still holds so that their errors + can be observed before they would otherwise be left orphaned. + """ + awaitables: list[Awaitable[Any]] = [] + with suppress_exceptions: + awaitables.extend(item for item in iterator if is_awaitable(item)) + return awaitables + + +def get_typename(value: Any) -> str | None: + """Get the ``__typename`` property of the given value.""" + if isinstance(value, Mapping): + return value.get("__typename") + # need to de-mangle the attribute assumed to be "private" in Python + for cls in value.__class__.__mro__: + __typename = getattr(value, f"_{cls.__name__}__typename", None) + if __typename: + return __typename + return None + + +def default_type_resolver( + value: Any, info: GraphQLResolveInfo, abstract_type: GraphQLAbstractType +) -> AwaitableOrValue[str | None]: + """Default type resolver function. + + If a resolve_type function is not given, then a default resolve behavior is used + which attempts two strategies: + + First, See if the provided value has a ``__typename`` field defined, if so, use that + value as name of the resolved type. + + Otherwise, test each possible type for the abstract type by calling + :meth:`~graphql.type.GraphQLObjectType.is_type_of` for the object + being coerced, returning the first type that matches. + """ + # First, look for `__typename`. + type_name = get_typename(value) + if isinstance(type_name, str): + return type_name + + # Otherwise, test each possible type. + possible_types = info.schema.get_possible_types(abstract_type) + is_awaitable = info.is_awaitable + awaitable_is_type_of_results: list[Awaitable[bool]] = [] + append_awaitable_result = awaitable_is_type_of_results.append + awaitable_types: list[GraphQLObjectType] = [] + append_awaitable_type = awaitable_types.append + + try: + for type_ in possible_types: + if type_.is_type_of: + is_type_of_result = type_.is_type_of(value, info) + + if is_awaitable(is_type_of_result): + append_awaitable_result(cast("Awaitable[bool]", is_type_of_result)) + append_awaitable_type(type_) + elif is_type_of_result: + if awaitable_is_type_of_results: + info.async_helpers.track(awaitable_is_type_of_results) + return type_.name + except Exception: + if awaitable_is_type_of_results: + # Settle the pending isTypeOf results in the background so that + # their errors can be observed before they would be orphaned. + info.async_helpers.track(awaitable_is_type_of_results) + raise + + if awaitable_is_type_of_results: + + async def get_type() -> str | None: + is_type_of_results = await info.async_helpers.gather( + awaitable_is_type_of_results + ) + for is_type_of_result, type_ in zip( + is_type_of_results, awaitable_types, strict=True + ): + if is_type_of_result: + return type_.name + return None + + return get_type() + + return None + + +def default_field_resolver(source: Any, info: GraphQLResolveInfo, **args: Any) -> Any: + """Default field resolver. + + If a resolve function is not given, then a default resolve behavior is used which + takes the property of the source object of the same name as the field and returns + it as the result, or if it's a function, returns the result of calling that function + while passing along args and context. + + For dictionaries, the field names are used as keys, for all other objects they are + used as attribute names. + """ + # Ensure source is a value for which property access is acceptable. + field_name = info.field_name + value = ( + source.get(field_name) + if isinstance(source, Mapping) + else getattr(source, field_name, None) + ) + if callable(value): + return value(info, **args) + return value diff --git a/src/graphql/execution/executor_throwing_on_incremental.py b/src/graphql/execution/executor_throwing_on_incremental.py new file mode 100644 index 00000000..940185ed --- /dev/null +++ b/src/graphql/execution/executor_throwing_on_incremental.py @@ -0,0 +1,104 @@ +"""Executor throwing on incremental delivery""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from ..error import GraphQLError +from ..language import OperationType +from .executor import ( + DEFER_NOT_SUPPORTED_ON_SUBSCRIPTIONS, + UNEXPECTED_MULTIPLE_PAYLOADS, + Executor, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterable, Iterable, Sequence + + from ..pyutils import AwaitableOrValue, Path + from ..type import ( + GraphQLList, + GraphQLObjectType, + GraphQLOutputType, + GraphQLResolveInfo, + ) + from .collect_fields import DeferUsage, FieldDetailsList, GroupedFieldSet + +__all__ = [ + "ExecutorThrowingOnIncremental", +] + + +class ExecutorThrowingOnIncremental(Executor[None]): + """Executor raising an error when the operation would defer or stream. + + For internal use only. + """ + + def execute_collected_root_fields( + self, + root_type: GraphQLObjectType, + root_value: Any, + grouped_field_set: GroupedFieldSet, + serially: bool, + new_defer_usages: Sequence[DeferUsage], + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the collected root fields, raising if something is deferred.""" + if new_defer_usages: + if self.operation.operation == OperationType.SUBSCRIPTION: + raise GraphQLError(DEFER_NOT_SUPPORTED_ON_SUBSCRIPTIONS) + raise GraphQLError(UNEXPECTED_MULTIPLE_PAYLOADS) + return self.execute_root_grouped_field_set( + root_type, + root_value, + grouped_field_set, + serially, + None, + ) + + def execute_collected_subfields( + self, + parent_type: GraphQLObjectType, + source_value: Any, + path: Path, + grouped_field_set: GroupedFieldSet, + new_defer_usages: Sequence[DeferUsage], + position_context: None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the collected subfields, raising if something is deferred.""" + if new_defer_usages: + if self.operation.operation == OperationType.SUBSCRIPTION: + raise GraphQLError(DEFER_NOT_SUPPORTED_ON_SUBSCRIPTIONS) + raise GraphQLError(UNEXPECTED_MULTIPLE_PAYLOADS) + + return self.execute_fields( + parent_type, + source_value, + path, + grouped_field_set, + None, + ) + + def complete_list_value( + self, + return_type: GraphQLList[GraphQLOutputType], + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + path: Path, + result: AsyncIterable[Any] | Iterable[Any], + position_context: None, + ) -> AwaitableOrValue[list[Any]]: + """Complete a list value, raising if something is streamed.""" + # this raises located errors when `@stream` is used on subscriptions + stream_usage = self.get_stream_usage(field_details_list, path) + if stream_usage is not None: + raise GraphQLError(UNEXPECTED_MULTIPLE_PAYLOADS) + + return super().complete_list_value( + return_type, + field_details_list, + info, + path, + result, + position_context, + ) diff --git a/src/graphql/execution/incremental/__init__.py b/src/graphql/execution/incremental/__init__.py new file mode 100644 index 00000000..74910714 --- /dev/null +++ b/src/graphql/execution/incremental/__init__.py @@ -0,0 +1,78 @@ +"""Incremental delivery execution + +The :mod:`graphql.execution.incremental` package contains the incremental +delivery engine that schedules and publishes deferred and streamed payloads. +For internal use only. +""" + +from .build_execution_plan import ( + DeferUsageSet, + ExecutionPlan, + build_execution_plan, +) +from .computation import Computation +from .incremental_executor import ( + DeliveryGroup, + ExecutionGroup, + ExecutionGroupValue, + IncrementalExecutor, + ItemStream, + StreamItemValue, +) +from .incremental_publisher import ( + IncrementalPublisher, + IncrementalPublisherContext, +) +from .stream_item_queue import StreamItemQueue +from .work_queue import ( + Group, + GroupFailureEvent, + GroupSuccessEvent, + GroupValuesEvent, + Stream, + StreamFailureEvent, + StreamItem, + StreamQueue, + StreamSuccessEvent, + StreamValuesEvent, + TaskResult, + Work, + WorkQueue, + WorkQueueEvent, + WorkQueueTerminationEvent, + WorkResult, + WorkTask, +) + +__all__ = [ + "Computation", + "DeferUsageSet", + "DeliveryGroup", + "ExecutionGroup", + "ExecutionGroupValue", + "ExecutionPlan", + "Group", + "GroupFailureEvent", + "GroupSuccessEvent", + "GroupValuesEvent", + "IncrementalExecutor", + "IncrementalPublisher", + "IncrementalPublisherContext", + "ItemStream", + "Stream", + "StreamFailureEvent", + "StreamItem", + "StreamItemQueue", + "StreamItemValue", + "StreamQueue", + "StreamSuccessEvent", + "StreamValuesEvent", + "TaskResult", + "Work", + "WorkQueue", + "WorkQueueEvent", + "WorkQueueTerminationEvent", + "WorkResult", + "WorkTask", + "build_execution_plan", +] diff --git a/src/graphql/execution/build_execution_plan.py b/src/graphql/execution/incremental/build_execution_plan.py similarity index 93% rename from src/graphql/execution/build_execution_plan.py rename to src/graphql/execution/incremental/build_execution_plan.py index bb94df16..c5429f73 100644 --- a/src/graphql/execution/build_execution_plan.py +++ b/src/graphql/execution/incremental/build_execution_plan.py @@ -1,82 +1,82 @@ -"""Build field plan""" - -from __future__ import annotations - -from typing import NamedTuple, TypeAlias - -from ..pyutils import RefMap, RefSet -from .collect_fields import DeferUsage, FieldDetailsList, GroupedFieldSet - -__all__ = [ - "DeferUsageSet", - "ExecutionPlan", - "GroupedFieldSet", - "build_execution_plan", -] - - -DeferUsageSet: TypeAlias = RefSet[DeferUsage] - - -class ExecutionPlan(NamedTuple): - """A plan for executing fields.""" - - grouped_field_set: GroupedFieldSet - new_grouped_field_sets: RefMap[DeferUsageSet, GroupedFieldSet] - - -def build_execution_plan( - original_grouped_field_set: GroupedFieldSet, - parent_defer_usages: DeferUsageSet | None = None, -) -> ExecutionPlan: - """Build a plan for executing fields.""" - if parent_defer_usages is None: - parent_defer_usages = RefSet() - - grouped_field_set: GroupedFieldSet = {} - new_grouped_field_sets: RefMap[DeferUsageSet, GroupedFieldSet] = RefMap() - - for response_key, field_details_list in original_grouped_field_set.items(): - filtered_defer_usage_set = get_filtered_defer_usage_set(field_details_list) - - if filtered_defer_usage_set == parent_defer_usages: - grouped_field_set[response_key] = field_details_list - continue - - for defer_usage_set in new_grouped_field_sets: - if defer_usage_set == filtered_defer_usage_set: - new_grouped_field_set = new_grouped_field_sets[defer_usage_set] - break - else: - new_grouped_field_set = {} - new_grouped_field_sets[filtered_defer_usage_set] = new_grouped_field_set - - new_grouped_field_set[response_key] = field_details_list - - return ExecutionPlan(grouped_field_set, new_grouped_field_sets) - - -def get_filtered_defer_usage_set( - field_details_list: FieldDetailsList, -) -> RefSet[DeferUsage]: - """Get a filtered set of defer usages.""" - # Create the set of defer usages for the field group. - filtered_defer_usage_set: RefSet[DeferUsage] = RefSet() - for field_details in field_details_list: - defer_usage = field_details.defer_usage - if defer_usage is None: - filtered_defer_usage_set.clear() - return filtered_defer_usage_set - filtered_defer_usage_set.add(defer_usage) - - # Remove defer usages that have a parent defer usage in the set. - # Since we remove in place, we need to iterate over a copy of the set. - for defer_usage in tuple(filtered_defer_usage_set): - parent_defer_usage: DeferUsage | None = defer_usage.parent_defer_usage - while parent_defer_usage is not None: - if parent_defer_usage in filtered_defer_usage_set: - filtered_defer_usage_set.discard(defer_usage) - break - parent_defer_usage = parent_defer_usage.parent_defer_usage - - return filtered_defer_usage_set +"""Build field plan""" + +from __future__ import annotations + +from typing import NamedTuple, TypeAlias + +from ...pyutils import RefMap, RefSet +from ..collect_fields import DeferUsage, FieldDetailsList, GroupedFieldSet + +__all__ = [ + "DeferUsageSet", + "ExecutionPlan", + "GroupedFieldSet", + "build_execution_plan", +] + + +DeferUsageSet: TypeAlias = RefSet[DeferUsage] + + +class ExecutionPlan(NamedTuple): + """A plan for executing fields.""" + + grouped_field_set: GroupedFieldSet + new_grouped_field_sets: RefMap[DeferUsageSet, GroupedFieldSet] + + +def build_execution_plan( + original_grouped_field_set: GroupedFieldSet, + parent_defer_usages: DeferUsageSet | None = None, +) -> ExecutionPlan: + """Build a plan for executing fields.""" + if parent_defer_usages is None: + parent_defer_usages = RefSet() + + grouped_field_set: GroupedFieldSet = {} + new_grouped_field_sets: RefMap[DeferUsageSet, GroupedFieldSet] = RefMap() + + for response_key, field_details_list in original_grouped_field_set.items(): + filtered_defer_usage_set = get_filtered_defer_usage_set(field_details_list) + + if filtered_defer_usage_set == parent_defer_usages: + grouped_field_set[response_key] = field_details_list + continue + + for defer_usage_set in new_grouped_field_sets: + if defer_usage_set == filtered_defer_usage_set: + new_grouped_field_set = new_grouped_field_sets[defer_usage_set] + break + else: + new_grouped_field_set = {} + new_grouped_field_sets[filtered_defer_usage_set] = new_grouped_field_set + + new_grouped_field_set[response_key] = field_details_list + + return ExecutionPlan(grouped_field_set, new_grouped_field_sets) + + +def get_filtered_defer_usage_set( + field_details_list: FieldDetailsList, +) -> RefSet[DeferUsage]: + """Get a filtered set of defer usages.""" + # Create the set of defer usages for the field group. + filtered_defer_usage_set: RefSet[DeferUsage] = RefSet() + for field_details in field_details_list: + defer_usage = field_details.defer_usage + if defer_usage is None: + filtered_defer_usage_set.clear() + return filtered_defer_usage_set + filtered_defer_usage_set.add(defer_usage) + + # Remove defer usages that have a parent defer usage in the set. + # Since we remove in place, we need to iterate over a copy of the set. + for defer_usage in tuple(filtered_defer_usage_set): + parent_defer_usage: DeferUsage | None = defer_usage.parent_defer_usage + while parent_defer_usage is not None: + if parent_defer_usage in filtered_defer_usage_set: + filtered_defer_usage_set.discard(defer_usage) + break + parent_defer_usage = parent_defer_usage.parent_defer_usage + + return filtered_defer_usage_set diff --git a/src/graphql/execution/incremental/computation.py b/src/graphql/execution/incremental/computation.py new file mode 100644 index 00000000..74a51d1d --- /dev/null +++ b/src/graphql/execution/incremental/computation.py @@ -0,0 +1,137 @@ +"""Memoized abortable computation""" + +from __future__ import annotations + +from asyncio import CancelledError, ensure_future +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +from ...pyutils.is_awaitable import is_awaitable + +if TYPE_CHECKING: + from asyncio import Future + from collections.abc import Callable + + from ...pyutils import AwaitableOrValue + +__all__ = ["Computation"] + +T = TypeVar("T") + +_PENDING = "pending" +_FULFILLED = "fulfilled" +_REJECTED = "rejected" + + +class Computation(Generic[T]): + """A lazily primed, memoized, abortable computation. + + Runs the given function at most once, no matter how often the computation + is primed or its result is requested. The outcome is memoized, whether the + function returns a value, returns an awaitable, or raises. An awaitable + result is scheduled as a future immediately, so that it can start running + early and can be cancelled on abort. + + Aborting an unprimed computation poisons it so that the function never + runs. Aborting a computation with a still pending future cancels the + future, invokes the optional abort callback and returns its result. + Aborting a settled computation has no effect. Requesting the result of a + computation that was aborted without a reason raises a CancelledError. + + For internal use only. + """ + + __slots__ = "_fn", "_on_abort", "_status", "_value" + + _status: str | None + _value: Any + + def __init__( + self, + fn: Callable[[], AwaitableOrValue[T]], + on_abort: Callable[[BaseException | None], AwaitableOrValue[None]] + | None = None, + ) -> None: + """Initialize the computation with a function and an abort callback.""" + self._fn = fn + self._on_abort = on_abort + self._status = None + self._value = None + + def prime(self) -> None: + """Run the computation if it has not been started or aborted yet.""" + if self._status is not None: + return + try: + result = self._fn() + except Exception as reason: + self._status = _REJECTED + self._value = reason + return + if is_awaitable(result): + future = ensure_future(result) + self._status = _PENDING + self._value = future + future.add_done_callback(self._settle) + else: + self._status = _FULFILLED + self._value = result + + @property + def pending_future(self) -> Future[T] | None: + """Get the still pending future, or None if there is none.""" + return self._value if self._status is _PENDING else None + + def result(self) -> AwaitableOrValue[T]: + """Get the memoized result, priming the computation if necessary. + + Returns the result or a future for it, or raises the reason why the + computation failed or was aborted. + """ + self.prime() + status = self._status + if status is _FULFILLED: + return self._value + if status is _REJECTED: + reason = self._value + if reason is None: + raise CancelledError + raise reason + return self._value # the pending future + + def abort(self, reason: BaseException | None = None) -> AwaitableOrValue[None]: + """Abort the computation with an optional reason. + + Returns the result of the abort callback when the computation was + still pending, None otherwise. + """ + status = self._status + if status is None: + self._status = _REJECTED + self._value = reason + elif status is _PENDING: + future = self._value + self._status = _REJECTED + self._value = reason + future.cancel() + on_abort = self._on_abort + if on_abort is not None: + return on_abort(reason) + return None + + def _settle(self, future: Future[T]) -> None: + """Memoize the outcome when the scheduled future is done.""" + if future.cancelled(): + if self._status is _PENDING: + # cancelled externally, not via abort + self._status = _REJECTED + self._value = CancelledError() + return + reason = future.exception() # always retrieve, avoiding warnings + if self._status is not _PENDING: + return # aborted after the future had already completed + if reason is not None: + self._status = _REJECTED + self._value = reason + else: + self._status = _FULFILLED + self._value = future.result() diff --git a/src/graphql/execution/incremental/incremental_executor.py b/src/graphql/execution/incremental/incremental_executor.py new file mode 100644 index 00000000..f1c3fdcd --- /dev/null +++ b/src/graphql/execution/incremental/incremental_executor.py @@ -0,0 +1,802 @@ +"""Incremental GraphQL executor""" + +from __future__ import annotations + +from asyncio import ensure_future, gather +from contextlib import suppress +from copy import copy +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +from ...error import GraphQLError, located_error +from ...language import OperationType +from ...pyutils import RefMap +from ..executor import ( + DEFER_NOT_SUPPORTED_ON_SUBSCRIPTIONS, + CollectedErrors, + Executor, + collect_iterator_awaitables, + to_nodes, +) +from .build_execution_plan import build_execution_plan +from .computation import Computation +from .incremental_publisher import IncrementalPublisher +from .stream_item_queue import StreamItemQueue +from .work_queue import Work, WorkResult, WorkTask + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterator, Sequence + + from ...pyutils import AwaitableOrValue, Path + from ...type import ( + GraphQLObjectType, + GraphQLOutputType, + GraphQLResolveInfo, + ) + from ..collect_fields import ( + DeferUsage, + FieldDetailsList, + GroupedFieldSet, + ) + from ..executor import StreamUsage + from ..types import ExecutionResult, ExperimentalIncrementalExecutionResults + from .build_execution_plan import DeferUsageSet, ExecutionPlan + from .work_queue import StreamQueue + +__all__ = [ + "DeliveryGroup", + "DeliveryGroupMap", + "ExecutionGroup", + "ExecutionGroupValue", + "IncrementalExecutor", + "ItemStream", + "StreamItemValue", + "should_defer", +] + +suppress_exceptions = suppress(Exception) + + +class DeliveryGroup: + """A deferred fragment addressable as a unit of delivery. + + Compared and hashed by identity, like all graph nodes. + + For internal use only. + """ + + __slots__ = "label", "parent", "path" + + path: Path | None + label: str | None + parent: DeliveryGroup | None + + def __init__( + self, + path: Path | None, + label: str | None, + parent: DeliveryGroup | None, + ) -> None: + self.path = path + self.label = label + self.parent = parent + + def __repr__(self) -> str: + name = self.__class__.__name__ + args: list[str] = [] + if self.path: + args.append(f"path={self.path.as_list()!r}") + if self.label: + args.append(f"label={self.label!r}") + if self.parent: + args.append("parent") + return f"{name}({', '.join(args)})" + + +DeliveryGroupMap = RefMap["DeferUsage", DeliveryGroup] + + +class ExecutionGroup(WorkTask): + """A deferred grouped field set belonging to one or more delivery groups. + + For internal use only. + """ + + __slots__ = ("path",) + + path: Path | None + + def __init__( + self, + groups: Sequence[DeliveryGroup], + computation: Computation[WorkResult], + path: Path | None, + ) -> None: + super().__init__(groups, computation) + self.path = path + + +class ItemStream: + """A streamed list field backed by a queue of stream item results. + + Compared and hashed by identity, like all graph nodes. + + For internal use only. + """ + + __slots__ = "initial_count", "label", "path", "queue" + + path: Path + label: str | None + queue: StreamQueue # a StreamItemQueue in practice + initial_count: int + + def __init__( + self, + path: Path, + label: str | None, + queue: StreamItemQueue, + initial_count: int, + ) -> None: + self.path = path + self.label = label + self.queue = queue + self.initial_count = initial_count + + def __repr__(self) -> str: + name = self.__class__.__name__ + args: list[str] = [f"path={self.path.as_list()!r}"] + if self.label: + args.append(f"label={self.label!r}") + return f"{name}({', '.join(args)})" + + +class ExecutionGroupValue(NamedTuple): + """The value produced by a completed execution group.""" + + delivery_groups: Sequence[DeliveryGroup] + path: list[str | int] + data: dict[str, Any] + errors: list[GraphQLError] | None = None + + +class StreamItemValue(NamedTuple): + """The value produced by a completed stream item.""" + + item: Any + errors: list[GraphQLError] | None = None + + +class IncrementalExecutor(Executor[DeliveryGroupMap]): + """Executor supporting incremental delivery via ``@defer`` and ``@stream``. + + Produces incremental work as data: each deferred grouped field set becomes + an execution group wrapping its executable in a computation, and each + streamed list field becomes an item stream owning a stream item queue. + Each execution group and stream item is executed by a fresh sub-executor + with its own collected errors, sharing the overall execution state. + + For internal use only. + """ + + defer_usage_set: DeferUsageSet | None + groups: list[DeliveryGroup] + tasks: list[ExecutionGroup] + streams: list[ItemStream] + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.defer_usage_set = None + self.groups = [] + self.tasks = [] + self.streams = [] + # Execution plan memoization shared with all sub-executors + self._initial_execution_plans: RefMap[GroupedFieldSet, ExecutionPlan] = RefMap() + self._deferred_execution_plans: RefMap[ + GroupedFieldSet, list[tuple[DeferUsageSet, ExecutionPlan]] + ] = RefMap() + + def create_sub_executor( + self, defer_usage_set: DeferUsageSet | None = None + ) -> IncrementalExecutor: + """Create a sub-executor with fresh per-execution state. + + The sub-executor is a shallow copy sharing the overall execution + state, with its own collected errors and produced incremental work. + """ + sub_executor = copy(self) + sub_executor.defer_usage_set = defer_usage_set + sub_executor.collected_errors = CollectedErrors() + sub_executor.groups = [] + sub_executor.tasks = [] + sub_executor.streams = [] + return sub_executor + + def abort(self, reason: BaseException | None = None) -> AwaitableOrValue[None]: + """Abort the incremental work produced by this executor. + + Aborts the computations of all execution groups and the queues of all + item streams produced by this executor, returning an awaitable for the + asynchronous part of the cleanup, or None when the whole cleanup could + be run synchronously. + """ + awaitables: list[Any] = [] + is_awaitable = self.is_awaitable + for task in self.tasks: + abort_result = task.computation.abort(reason) + if is_awaitable(abort_result): + awaitables.append(abort_result) + for stream in self.streams: + abort_result = stream.queue.abort(reason) + if is_awaitable(abort_result): + awaitables.append(abort_result) + if not awaitables: + return None + + async def settle_awaitables() -> None: + await gather(*awaitables, return_exceptions=True) + + return settle_awaitables() + + def abort_in_background(self, reason: BaseException | None = None) -> None: + """Abort the produced incremental work, settling cleanup in background. + + Used on synchronous failure paths that cannot await the asynchronous + part of the cleanup. + """ + abort_result = self.abort(reason) + if self.is_awaitable(abort_result): + self.settle_in_background([abort_result]) + + def build_response( + self, data: dict[str, Any] | None + ) -> ExecutionResult | ExperimentalIncrementalExecutionResults: + """Build the response for the given completed initial data. + + When execution groups or item streams are still pending, an + incremental response is built which delivers them via an incremental + publisher; otherwise, the normal response is built. + """ + work = self.get_incremental_work() + if not work.tasks and not work.streams: + return super().build_response(data) + + errors = self.collected_errors.errors + incremental_publisher = IncrementalPublisher() + return incremental_publisher.build_response( + cast("dict[str, Any]", data), list(errors) or None, work, self + ) + + def execute_collected_root_fields( + self, + root_type: GraphQLObjectType, + root_value: Any, + grouped_field_set: GroupedFieldSet, + serially: bool, + new_defer_usages: Sequence[DeferUsage], + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the collected root fields with incremental delivery.""" + if not new_defer_usages: + return self.execute_root_grouped_field_set( + root_type, + root_value, + grouped_field_set, + serially, + None, + ) + + if self.operation.operation == OperationType.SUBSCRIPTION: + raise GraphQLError(DEFER_NOT_SUPPORTED_ON_SUBSCRIPTIONS) + + new_delivery_groups, new_delivery_group_map = self.get_new_delivery_group_map( + new_defer_usages, None, None + ) + + execution_plan = self.build_root_execution_plan(grouped_field_set) + planned_field_set, new_grouped_field_sets = execution_plan + + data = self.execute_root_grouped_field_set( + root_type, + root_value, + planned_field_set, + serially, + new_delivery_group_map, + ) + + self.groups.extend(new_delivery_groups) + + if new_grouped_field_sets: + self.collect_execution_groups( + root_type, + root_value, + None, + new_grouped_field_sets, + new_delivery_group_map, + ) + + return data + + def build_root_execution_plan( + self, original_grouped_field_set: GroupedFieldSet + ) -> ExecutionPlan: + """Build a memoized execution plan for the root fields.""" + execution_plans = self._initial_execution_plans + execution_plan = execution_plans.get(original_grouped_field_set) + if execution_plan is None: + execution_plan = build_execution_plan(original_grouped_field_set) + execution_plans[original_grouped_field_set] = execution_plan + return execution_plan + + def execute_collected_subfields( + self, + parent_type: GraphQLObjectType, + source_value: Any, + path: Path, + grouped_field_set: GroupedFieldSet, + new_defer_usages: Sequence[DeferUsage], + delivery_group_map: DeliveryGroupMap | None, + ) -> AwaitableOrValue[dict[str, Any]]: + """Execute the collected subfields with incremental delivery.""" + if new_defer_usages and self.operation.operation == OperationType.SUBSCRIPTION: + raise GraphQLError(DEFER_NOT_SUPPORTED_ON_SUBSCRIPTIONS) + + if delivery_group_map is None and not new_defer_usages: + return self.execute_fields( + parent_type, + source_value, + path, + grouped_field_set, + delivery_group_map, + ) + + new_delivery_groups, new_delivery_group_map = self.get_new_delivery_group_map( + new_defer_usages, delivery_group_map, path + ) + + execution_plan = self.build_sub_execution_plan(grouped_field_set) + planned_field_set, new_grouped_field_sets = execution_plan + + data = self.execute_fields( + parent_type, + source_value, + path, + planned_field_set, + new_delivery_group_map, + ) + + self.groups.extend(new_delivery_groups) + + if new_grouped_field_sets: + self.collect_execution_groups( + parent_type, + source_value, + path, + new_grouped_field_sets, + new_delivery_group_map, + ) + + return data + + def build_sub_execution_plan( + self, original_grouped_field_set: GroupedFieldSet + ) -> ExecutionPlan: + """Build a memoized execution plan for the subfields.""" + defer_usage_set = self.defer_usage_set + if defer_usage_set is None: + return self.build_root_execution_plan(original_grouped_field_set) + execution_plans = self._deferred_execution_plans + planned = execution_plans.get(original_grouped_field_set) + if planned is None: + planned = [] + execution_plans[original_grouped_field_set] = planned + else: + for planned_defer_usage_set, execution_plan in planned: + if planned_defer_usage_set is defer_usage_set: + return execution_plan + execution_plan = build_execution_plan( + original_grouped_field_set, defer_usage_set + ) + planned.append((defer_usage_set, execution_plan)) + return execution_plan + + def collect_execution_groups( + self, + parent_type: GraphQLObjectType, + source_value: Any, + path: Path | None, + new_grouped_field_sets: RefMap[DeferUsageSet, GroupedFieldSet], + delivery_group_map: DeliveryGroupMap, + ) -> None: + """Create execution groups for the new deferred grouped field sets.""" + append_task = self.tasks.append + parent_defer_usages = self.defer_usage_set + enable_early_execution = self.enable_early_execution + for defer_usage_set, grouped_field_set in new_grouped_field_sets.items(): + delivery_groups = get_delivery_groups(defer_usage_set, delivery_group_map) + + sub_executor = self.create_sub_executor(defer_usage_set) + + def execute_group( + sub_executor: IncrementalExecutor = sub_executor, + delivery_groups: Sequence[DeliveryGroup] = delivery_groups, + grouped_field_set: GroupedFieldSet = grouped_field_set, + ) -> AwaitableOrValue[WorkResult]: + return sub_executor.execute_execution_group( + delivery_groups, + parent_type, + source_value, + path, + grouped_field_set, + delivery_group_map, + ) + + computation = Computation(execute_group, sub_executor.abort) + + if enable_early_execution: + if should_defer(parent_defer_usages, defer_usage_set): + self.prime_soon(computation) + else: + self.prime_now(computation) + + append_task(ExecutionGroup(delivery_groups, computation, path)) + + def prime_now(self, computation: Computation[WorkResult]) -> None: + """Prime the computation immediately, tracking its pending future.""" + computation.prime() + future = computation.pending_future + if future is not None: + self.track_incremental_future(future) + + def prime_soon(self, computation: Computation[WorkResult]) -> None: + """Prime the computation early, but only after the current work. + + This makes sure that a new deferred execution group does not run + before the execution step that created it has finished. + """ + + async def prime() -> None: + self.prime_now(computation) + + self.track_incremental_future(ensure_future(prime())) + + def execute_execution_group( + self, + delivery_groups: Sequence[DeliveryGroup], + parent_type: GraphQLObjectType, + source_value: Any, + path: Path | None, + grouped_field_set: GroupedFieldSet, + delivery_group_map: DeliveryGroupMap, + ) -> AwaitableOrValue[WorkResult]: + """Execute an execution group on this sub-executor.""" + try: + result = self.execute_fields( + parent_type, + source_value, + path, + grouped_field_set, + delivery_group_map, + ) + except Exception: + self.abort_in_background() + raise + + if self.is_awaitable(result): + + async def await_result() -> WorkResult: + try: + data = await result + except Exception: + abort_result = self.abort() + if self.is_awaitable(abort_result): + await abort_result + raise + return self.build_execution_group_result(delivery_groups, path, data) + + return await_result() + + return self.build_execution_group_result( + delivery_groups, path, cast("dict[str, Any]", result) + ) + + def build_execution_group_result( + self, + delivery_groups: Sequence[DeliveryGroup], + path: Path | None, + data: dict[str, Any], + ) -> WorkResult: + """Build the result of a completed execution group.""" + errors = self.collected_errors.errors + value = ExecutionGroupValue( + delivery_groups, + path.as_list() if path else [], + data, + list(errors) if errors else None, + ) + return WorkResult(value, self.get_incremental_work()) + + def get_incremental_work(self) -> Work: + """Get the incremental work produced by this executor. + + When errors have been collected, execution groups and item streams + whose position has been nulled via error propagation are aborted and + filtered out before they reach the scheduler, cancelling work that is + no longer being delivered. + """ + groups, tasks, streams = self.groups, self.tasks, self.streams + collected_errors = self.collected_errors + + if not collected_errors.errors: + return Work(groups, tasks, streams) + + has_nulled_position = collected_errors.has_nulled_position + cancellation_reason = RuntimeError( + "Cancelled secondary to null within original result" + ) + + filtered_tasks: list[ExecutionGroup] = [] + for task in tasks: + if has_nulled_position(task.path): + self.settle_abort_result(task.computation.abort(cancellation_reason)) + else: + filtered_tasks.append(task) + + filtered_streams: list[ItemStream] = [] + for stream in streams: + if has_nulled_position(stream.path): + self.settle_abort_result(stream.queue.abort(cancellation_reason)) + else: + filtered_streams.append(stream) + + return Work(groups, filtered_tasks, filtered_streams) + + def settle_abort_result(self, abort_result: AwaitableOrValue[None]) -> None: + """Settle the asynchronous part of an abort in the background.""" + if self.is_awaitable(abort_result): + self.settle_in_background([abort_result]) + + def get_new_delivery_group_map( + self, + new_defer_usages: Sequence[DeferUsage], + delivery_group_map: DeliveryGroupMap | None, + path: Path | None, + ) -> tuple[list[DeliveryGroup], DeliveryGroupMap]: + """Get the new delivery group map. + + Instantiates new DeliveryGroups for the given path, returning an + updated map of DeferUsage objects to DeliveryGroups. + + Note: As defer directives may be used with operations returning lists, + a DeferUsage object may correspond to many DeliveryGroups. + """ + new_delivery_groups: list[DeliveryGroup] = [] + new_delivery_group_map: DeliveryGroupMap = RefMap( + None if delivery_group_map is None else delivery_group_map.items() + ) + + # For each new DeferUsage object: + for new_defer_usage in new_defer_usages: + parent_defer_usage = new_defer_usage.parent_defer_usage + + parent = ( + None + if parent_defer_usage is None + else delivery_group_from_defer_usage( + parent_defer_usage, new_delivery_group_map + ) + ) + + # Instantiate the new delivery group. + delivery_group = DeliveryGroup(path, new_defer_usage.label, parent) + + # Add it to the list of new delivery groups. + new_delivery_groups.append(delivery_group) + + # Update the map. + new_delivery_group_map[new_defer_usage] = delivery_group + + return new_delivery_groups, new_delivery_group_map + + def handle_stream( + self, + index: int, + path: Path, + iterator: Iterator[Any] | AsyncIterator[Any], + is_async: bool, + stream_usage: StreamUsage, + info: GraphQLResolveInfo, + item_type: GraphQLOutputType, + ) -> bool: + """Stream the remaining list items via an item stream.""" + queue = self.build_stream_item_queue( + index, + path, + iterator, + stream_usage.field_details_list, + info, + item_type, + is_async, + ) + + item_stream = ItemStream(path, stream_usage.label, queue, index) + + self.streams.append(item_stream) + return True + + def build_stream_item_queue( + self, + initial_index: int, + stream_path: Path, + iterator: Iterator[Any] | AsyncIterator[Any], + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + item_type: GraphQLOutputType, + is_async: bool, + ) -> StreamItemQueue: + """Build the stream item queue for a streamed list field.""" + enable_early_execution = self.enable_early_execution + is_awaitable = self.is_awaitable + + async def produce(queue: StreamItemQueue) -> None: + index = initial_index + while True: + try: + item = ( + await anext(cast("AsyncIterator[Any]", iterator)) + if is_async + else next(cast("Iterator[Any]", iterator)) + ) + except (StopAsyncIteration, StopIteration): + return + except Exception as raw_error: + raise located_error( + raw_error, to_nodes(field_details_list), stream_path.as_list() + ) from raw_error + + item_path = stream_path.add_key(index, None) + + sub_executor = self.create_sub_executor() + + result = sub_executor.complete_stream_item( + item_path, item, field_details_list, info, item_type + ) + if is_awaitable(result): + if enable_early_execution: + await queue.push(ensure_future(result)) + else: + await queue.push(await result) + else: + await queue.push(cast("WorkResult", result)) + + index += 1 + + def on_abort(_reason: BaseException | None) -> AwaitableOrValue[None]: + if is_async: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + + async def close_iterator() -> None: + with suppress_exceptions: + await aclose() + + return close_iterator() + return None + # Drain the sync iterator so that any awaitable items it still + # holds can be settled before they would be orphaned. + self.track_async_work( + collect_iterator_awaitables( + cast("Iterator[Any]", iterator), is_awaitable + ) + ) + return None + + return StreamItemQueue(produce, on_abort, eager=enable_early_execution) + + def complete_stream_item( + self, + item_path: Path, + item: Any, + field_details_list: FieldDetailsList, + info: GraphQLResolveInfo, + item_type: GraphQLOutputType, + ) -> AwaitableOrValue[WorkResult]: + """Complete a stream item on this sub-executor.""" + is_awaitable = self.is_awaitable + if is_awaitable(item): + + async def await_stream_item_result() -> WorkResult: + try: + completed = await self.complete_awaitable_value( + item_type, + field_details_list, + info, + item_path, + item, + None, + ) + except Exception: + abort_result = self.abort() + if is_awaitable(abort_result): + await abort_result + raise + return self.build_stream_item_result(completed) + + return await_stream_item_result() + + try: + try: + completed = self.complete_value( + item_type, + field_details_list, + info, + item_path, + item, + None, + ) + except Exception as raw_error: + self.handle_field_error( + raw_error, item_type, field_details_list, item_path + ) + return self.build_stream_item_result(None) + except Exception: + self.abort_in_background() + raise + + if is_awaitable(completed): + + async def await_completed_stream_item_result() -> WorkResult: + try: + try: + resolved = await completed + except Exception as raw_error: + self.handle_field_error( + raw_error, item_type, field_details_list, item_path + ) + resolved = None + except Exception: + abort_result = self.abort() + if is_awaitable(abort_result): + await abort_result + raise + return self.build_stream_item_result(resolved) + + return await_completed_stream_item_result() + + return self.build_stream_item_result(completed) + + def build_stream_item_result(self, item: Any) -> WorkResult: + """Build the result of a completed stream item.""" + errors = self.collected_errors.errors + value = StreamItemValue(item, list(errors) if errors else None) + return WorkResult(value, self.get_incremental_work()) + + +def should_defer( + parent_defer_usages: DeferUsageSet | None, defer_usage_set: DeferUsageSet +) -> bool: + """Decide whether to defer the given defer usage set. + + If we have a new child defer usage, defer. + Otherwise, this defer usage was already deferred when it was initially + encountered, and is now in the midst of executing early, so the new + deferred grouped fields set can be executed immediately. + """ + return parent_defer_usages is None or not all( + defer_usage in parent_defer_usages for defer_usage in defer_usage_set + ) + + +def get_delivery_groups( + defer_usage_set: DeferUsageSet, delivery_group_map: DeliveryGroupMap +) -> list[DeliveryGroup]: + """Get the delivery groups for the given defer usages.""" + return [ + delivery_group_from_defer_usage(defer_usage, delivery_group_map) + for defer_usage in defer_usage_set + ] + + +def delivery_group_from_defer_usage( + defer_usage: DeferUsage, delivery_group_map: DeliveryGroupMap +) -> DeliveryGroup: + """Get the delivery group mapped to the given defer usage.""" + return delivery_group_map[defer_usage] diff --git a/src/graphql/execution/incremental/incremental_publisher.py b/src/graphql/execution/incremental/incremental_publisher.py new file mode 100644 index 00000000..456c89d2 --- /dev/null +++ b/src/graphql/execution/incremental/incremental_publisher.py @@ -0,0 +1,316 @@ +"""Incremental publisher""" + +from __future__ import annotations + +from asyncio import FIRST_COMPLETED, ensure_future, wait +from typing import TYPE_CHECKING, Any, Protocol, cast + +from ...error import GraphQLError, located_error +from ..types import ( + CompletedResult, + ExperimentalIncrementalExecutionResults, + IncrementalDeferResult, + IncrementalStreamResult, + InitialIncrementalExecutionResult, + PendingResult, + SubsequentIncrementalExecutionResult, +) +from .work_queue import ( + GroupFailureEvent, + GroupSuccessEvent, + GroupValuesEvent, + StreamFailureEvent, + StreamSuccessEvent, + StreamValuesEvent, + WorkQueue, +) + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Sequence + + from ...pyutils import AbortSignal + from ..types import IncrementalResult + from .incremental_executor import ( + DeliveryGroup, + ExecutionGroupValue, + ItemStream, + StreamItemValue, + ) + from .work_queue import Work, WorkQueueEvent + +__all__ = [ + "IncrementalPublisher", + "IncrementalPublisherContext", +] + + +class IncrementalPublisherContext(Protocol): + """The context for incremental publishing.""" + + abort_signal: AbortSignal | None + + def abort_error(self) -> Exception: + """Return the exception to raise when execution has been aborted.""" + ... # pragma: no cover + + async def cancel_incremental_work( + self, reason: BaseException | None = None + ) -> None: + """Cancel all pending incremental work and close the stream sources.""" + ... # pragma: no cover + + def run_async_work_finished_hook(self) -> None: + """Run the hook signaling that all asynchronous work has finished.""" + ... # pragma: no cover + + +class _SubsequentResultContext: + """The context collecting the parts of a subsequent incremental result.""" + + __slots__ = "completed", "has_next", "incremental", "pending" + + def __init__(self) -> None: + self.pending: list[PendingResult] = [] + self.incremental: list[IncrementalResult] = [] + self.completed: list[CompletedResult] = [] + self.has_next = True + + +class IncrementalPublisher: + """Publish incremental results. + + This class is used to publish incremental results to the client, enabling + semi-concurrent execution while preserving result order. It translates the + work queue events into incremental execution results, assigning the ids of + the pending results. + + For internal use only. + """ + + _ids: dict[DeliveryGroup | ItemStream, str] + _next_id: int + + def __init__(self) -> None: + self._ids = {} + self._next_id = 0 + + def build_response( + self, + data: dict[str, Any], + errors: list[GraphQLError] | None, + work: Work, + context: IncrementalPublisherContext, + ) -> ExperimentalIncrementalExecutionResults: + """Build the initial response and the subsequent result stream.""" + work_queue = WorkQueue(work) + + pending = self._to_pending_results( + cast("Sequence[DeliveryGroup]", work_queue.initial_groups), + cast("Sequence[ItemStream]", work_queue.initial_streams), + ) + + initial_result = InitialIncrementalExecutionResult( + data, + errors, + pending=pending, + has_next=True, + ) + + return ExperimentalIncrementalExecutionResults( + initial_result, self._subscribe(work_queue, context) + ) + + def _ensure_id(self, node: DeliveryGroup | ItemStream) -> str: + """Get the id assigned to the given node, assigning one if needed.""" + id_ = self._ids.get(node) + if id_ is None: + id_ = str(self._next_id) + self._next_id += 1 + self._ids[node] = id_ + return id_ + + def _to_pending_results( + self, + new_groups: Sequence[DeliveryGroup], + new_streams: Sequence[ItemStream], + ) -> list[PendingResult]: + """Convert new delivery groups and streams to pending results.""" + pending_results: list[PendingResult] = [] + nodes: list[DeliveryGroup | ItemStream] = [*new_groups, *new_streams] + for node in nodes: + id_ = self._ensure_id(node) + path = node.path + pending_results.append( + PendingResult(id_, path.as_list() if path else [], node.label) + ) + return pending_results + + async def _subscribe( + self, + work_queue: WorkQueue, + context: IncrementalPublisherContext, + ) -> AsyncGenerator[SubsequentIncrementalExecutionResult, None]: + """Subscribe to the incremental results.""" + abort_signal = context.abort_signal + events = work_queue.events() + try: + while True: + if abort_signal is not None and abort_signal.aborted: + raise context.abort_error() + + next_batch = ensure_future(anext(events)) + if abort_signal is None: + try: + batch = await next_batch + except StopAsyncIteration: + return + else: + # reject the pending request when the operation is aborted + abort = ensure_future(abort_signal.wait()) + try: + await wait({next_batch, abort}, return_when=FIRST_COMPLETED) + finally: + if not abort.done(): + abort.cancel() + if abort_signal.aborted: + next_batch.cancel() + raise context.abort_error() + try: + batch = next_batch.result() + except StopAsyncIteration: + return + + subsequent_result = self._handle_batch(batch) + yield subsequent_result + + if not subsequent_result.has_next: + return + finally: + await work_queue.cancel() + await context.cancel_incremental_work() + context.run_async_work_finished_hook() + + def _handle_batch( + self, batch: Sequence[WorkQueueEvent] + ) -> SubsequentIncrementalExecutionResult: + """Translate a batch of work queue events into a subsequent result.""" + context = _SubsequentResultContext() + + for event in batch: + self._handle_work_queue_event(event, context) + + return SubsequentIncrementalExecutionResult( + has_next=context.has_next, + pending=context.pending or None, + incremental=context.incremental or None, + completed=context.completed or None, + ) + + def _handle_work_queue_event( + self, + event: WorkQueueEvent, + context: _SubsequentResultContext, + ) -> None: + """Translate a single work queue event.""" + if isinstance(event, GroupValuesEvent): + group = cast("DeliveryGroup", event.group) + id_ = self._ensure_id(group) + for value in cast("Sequence[ExecutionGroupValue]", event.values): + best_id, sub_path = self._get_best_id_and_sub_path(id_, group, value) + context.incremental.append( + IncrementalDeferResult( + data=value.data, + id=best_id, + sub_path=sub_path, + errors=value.errors, + ) + ) + elif isinstance(event, GroupSuccessEvent): + group = cast("DeliveryGroup", event.group) + context.completed.append(CompletedResult(self._ensure_id(group))) + del self._ids[group] + if event.new_groups or event.new_streams: + context.pending.extend( + self._to_pending_results( + cast("Sequence[DeliveryGroup]", event.new_groups), + cast("Sequence[ItemStream]", event.new_streams), + ) + ) + elif isinstance(event, GroupFailureEvent): + group = cast("DeliveryGroup", event.group) + context.completed.append( + CompletedResult( + self._ensure_id(group), [ensure_graphql_error(event.error)] + ) + ) + del self._ids[group] + elif isinstance(event, StreamValuesEvent): + stream = cast("ItemStream", event.stream) + id_ = self._ensure_id(stream) + items: list[Any] = [] + errors: list[GraphQLError] = [] + for item_value in cast("Sequence[StreamItemValue]", event.values): + items.append(item_value.item) + if item_value.errors: + errors.extend(item_value.errors) + context.incremental.append( + IncrementalStreamResult(items=items, id=id_, errors=errors or None) + ) + if event.new_groups or event.new_streams: + context.pending.extend( + self._to_pending_results( + cast("Sequence[DeliveryGroup]", event.new_groups), + cast("Sequence[ItemStream]", event.new_streams), + ) + ) + elif isinstance(event, StreamSuccessEvent): + stream = cast("ItemStream", event.stream) + context.completed.append(CompletedResult(self._ensure_id(stream))) + del self._ids[stream] + elif isinstance(event, StreamFailureEvent): + stream = cast("ItemStream", event.stream) + context.completed.append( + CompletedResult( + self._ensure_id(stream), [ensure_graphql_error(event.error)] + ) + ) + del self._ids[stream] + else: # WorkQueueTerminationEvent + context.has_next = False + + def _get_best_id_and_sub_path( + self, + initial_id: str, + initial_delivery_group: DeliveryGroup, + value: ExecutionGroupValue, + ) -> tuple[str, list[str | int] | None]: + """Get the best id and sub path for an execution group value.""" + path = initial_delivery_group.path + max_length = len(path.as_list()) if path else 0 + best_id = initial_id + + for delivery_group in value.delivery_groups: + if delivery_group is initial_delivery_group: + continue + id_ = self._ids.get(delivery_group) + if id_ is None: + continue + group_path = delivery_group.path + length = len(group_path.as_list()) if group_path else 0 + if length > max_length: + max_length = length + best_id = id_ + + sub_path = value.path[max_length:] + return best_id, sub_path or None + + +def ensure_graphql_error(error: BaseException) -> GraphQLError: + """Ensure that the given error is returned as a GraphQL error.""" + if isinstance(error, GraphQLError): + return error + if isinstance(error, Exception): + return located_error(error) + # BaseExceptions like CancelledError normally do not surface here, + # since cancellation stops the delivery of events. + return GraphQLError(str(error) or repr(error)) diff --git a/src/graphql/execution/incremental/stream_item_queue.py b/src/graphql/execution/incremental/stream_item_queue.py new file mode 100644 index 00000000..716c649f --- /dev/null +++ b/src/graphql/execution/incremental/stream_item_queue.py @@ -0,0 +1,255 @@ +"""Queue of stream item results""" + +from __future__ import annotations + +from asyncio import ( + Event, + Future, + Queue, + QueueEmpty, + ensure_future, + gather, + get_running_loop, + isfuture, +) +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +from ...pyutils.is_awaitable import is_awaitable + +if TYPE_CHECKING: + from asyncio import Task + from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Sequence + + from .work_queue import WorkResult + +__all__ = ["StreamItemQueue"] + +_END: Any = object() # sentinel marking the normal end of the stream + + +class _ErrorEntry(NamedTuple): + """An entry marking the failure of the stream.""" + + error: BaseException + + +class StreamItemQueue: + """An ordered queue of stream item results with back-pressure. + + The queue is fed by a producer coroutine which pushes the stream item + results in source order, either settled or as pending futures when items + are completed early. The queue buffers up to ``capacity`` entries; when + the buffer is full, pushing blocks, which paces the producer. + + The batches of settled results are consumed via :meth:`batches`, which + preserves the source order: a still pending head entry delays delivery, + while consecutive already settled entries are delivered together in one + batch. This implements the stream queue interface consumed by the + :class:`~graphql.execution.incremental.WorkQueue` scheduler. + + In eager mode, the producer is started as soon as possible; otherwise it + is started lazily when the batches are first consumed. + + Aborting the queue cancels the producer and all pending item futures and + awaits their cancellation before running the cleanup callback, so that no + asynchronous work is left behind unsettled. + + For internal use only. + """ + + def __init__( + self, + produce: Callable[[StreamItemQueue], Coroutine[Any, Any, None]], + on_abort: Callable[[BaseException | None], Awaitable[None] | None] + | None = None, + eager: bool = False, + capacity: int = 100, + ) -> None: + """Initialize the queue with a producer and an abort callback.""" + self._produce = produce + self._on_abort = on_abort + self._eager = eager + self._entries: Queue[Any] = Queue(capacity) + self._started = Event() + self._producer_task: Task[None] | None = None + self._producer_parked = False + self._producer_cancelled = False + self._pending_futures: set[Future[WorkResult]] = set() + self._aborted = False + self._finished = False + self._stopped = False + if eager: + try: + get_running_loop() + except RuntimeError: + pass # no running event loop yet, start on first consumption + else: + self._start() + + def _start(self) -> None: + """Start the producer task if it has not been started yet.""" + if self._producer_task is None and not self._aborted: + self._producer_task = ensure_future(self._run()) + + async def _run(self) -> None: + """Run the producer, finishing the queue when it is done.""" + if not self._eager: + await self._started.wait() + entries = self._entries + try: + await self._produce(self) + except Exception as error: + # settle the pending item futures and clean up the source + # before delivering the failure + self._aborted = True + await self._settle_pending() + on_abort = self._on_abort + if on_abort is not None: + cleanup = on_abort(error) + if is_awaitable(cleanup): + await cleanup + self._producer_parked = True # may park on the full queue + await entries.put(_ErrorEntry(error)) + else: + self._finished = True + self._producer_parked = True # may park on the full queue + await entries.put(_END) + + async def push(self, result: WorkResult | Future[WorkResult]) -> None: + """Push a settled or still pending stream item result onto the queue. + + Blocks when the queue has reached its capacity, pacing the producer. + """ + if isfuture(result): + pending_futures = self._pending_futures + pending_futures.add(result) + result.add_done_callback(pending_futures.discard) + await self._entries.put(result) + + async def batches(self) -> AsyncIterator[Sequence[WorkResult]]: + """Iterate over the batches of settled stream item results in order. + + Starts the producer on the first iteration step. A failure of the + stream is raised only after all results settled before the failure + have been delivered. + """ + self._start() + self._started.set() + entries = self._entries + held: Any = None + while True: + entry = await entries.get() if held is None else held + held = None + if isfuture(entry): + try: + entry = await entry + except Exception: + await self._cleanup() + raise + if entry is _END: + self._stopped = True + return + if isinstance(entry, _ErrorEntry): + raise entry.error + batch = [entry] + while True: + try: + next_entry = entries.get_nowait() + except QueueEmpty: + break + if next_entry is _END: + # allow peeking ahead to see that the stream has stopped + self._stopped = True + held = next_entry # finish after delivering this batch + break + if isinstance(next_entry, _ErrorEntry) or ( + isfuture(next_entry) and not next_entry.done() + ): + held = next_entry # deliver the current batch first + break + if isfuture(next_entry): + try: + next_entry = next_entry.result() + except Exception: + held = next_entry # re-raise when delivered as head + break + batch.append(next_entry) + yield batch + + def is_stopped(self) -> bool: + """Check whether the stream is known to have stopped normally.""" + return self._stopped + + def abort(self, reason: BaseException | None = None) -> Awaitable[None] | None: + """Abort the stream with an optional reason. + + Cancels the producer and the pending item futures and returns an + awaitable for the asynchronous part of the cleanup, or None when the + whole cleanup could be run synchronously. + """ + producer_task = self._producer_task + running = producer_task is not None and not producer_task.done() + parked = running and self._producer_parked and not self._producer_cancelled + if parked: + # release the producer parked on the back-pressured queue + # while trying to deliver its final entry + producer_task.cancel() # type: ignore[union-attr] + self._producer_cancelled = True + if self._aborted: + # Aborted (or failed) before, so the cleanup has already run; only + # release a producer that was still parked. + return self._settle_parked() if parked else None + self._aborted = True + if self._finished: + # The source finished normally, so it must not be cleaned up; only + # release a producer that was still parked and settle any still + # pending early executed item futures. + if not parked and not self._pending_futures: + return None + for future in self._pending_futures: + future.cancel() + return self._settle_parked() + if running and not self._producer_cancelled: + producer_task.cancel() # type: ignore[union-attr] + self._producer_cancelled = True + for future in self._pending_futures: + future.cancel() + if not running and not self._pending_futures: + # nothing to cancel asynchronously, just run the cleanup callback + on_abort = self._on_abort + if on_abort is not None: + cleanup = on_abort(reason) + if is_awaitable(cleanup): + return cleanup + return None + return self._cleanup(reason) + + async def _settle_parked(self) -> None: + """Await the cancelled parked producer and settle pending item futures.""" + # the callers guarantee that the producer task has been created + producer_task = cast("Task[None]", self._producer_task) + await gather(producer_task, return_exceptions=True) + await self._settle_pending() + + async def _settle_pending(self) -> None: + """Cancel and settle all still pending item futures.""" + pending = [future for future in self._pending_futures if not future.done()] + if pending: + for future in pending: + future.cancel() + await gather(*pending, return_exceptions=True) + + async def _cleanup(self, reason: BaseException | None = None) -> None: + """Cancel all pending work, awaiting it, and run the abort callback.""" + self._aborted = True + producer_task = self._producer_task + if producer_task is not None and not producer_task.done(): + producer_task.cancel() + self._producer_cancelled = True + await gather(producer_task, return_exceptions=True) + await self._settle_pending() + on_abort = self._on_abort + if on_abort is not None: + cleanup = on_abort(reason) + if is_awaitable(cleanup): + await cleanup diff --git a/src/graphql/execution/incremental/work_queue.py b/src/graphql/execution/incremental/work_queue.py new file mode 100644 index 00000000..93463bc6 --- /dev/null +++ b/src/graphql/execution/incremental/work_queue.py @@ -0,0 +1,669 @@ +"""Work queue scheduling incremental delivery""" + +from __future__ import annotations + +from asyncio import ( + CancelledError, + Event, + Queue, + QueueEmpty, + ensure_future, + gather, +) +from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, TypeAlias, cast + +from ...pyutils.is_awaitable import is_awaitable + +if TYPE_CHECKING: + from asyncio import Future, Task + from collections.abc import AsyncIterator, Awaitable, Sequence + + from .computation import Computation + +__all__ = [ + "Group", + "GroupFailureEvent", + "GroupSuccessEvent", + "GroupValuesEvent", + "Stream", + "StreamFailureEvent", + "StreamItem", + "StreamQueue", + "StreamSuccessEvent", + "StreamValuesEvent", + "TaskResult", + "Work", + "WorkQueue", + "WorkQueueEvent", + "WorkQueueTerminationEvent", + "WorkResult", + "WorkTask", +] + +_UNSET: Any = object() + +_STOP: Any = object() # sentinel waking up the event consumer on cancellation + + +class Group(Protocol): + """A group of work with an optional parent group. For internal use only.""" + + parent: Any # the parent Group or None + + +class StreamQueue(Protocol): + """The queue interface a stream must provide. For internal use only.""" + + def batches(self) -> AsyncIterator[Sequence[Any]]: + """Iterate over the batches of settled stream items.""" + ... + + def is_stopped(self) -> bool: + """Check whether the stream is known to have stopped normally.""" + ... + + def abort(self, reason: BaseException | None = None) -> Awaitable[None] | None: + """Abort the stream with an optional reason.""" + ... + + +class Stream(Protocol): + """A stream of work items backed by a queue. For internal use only.""" + + queue: StreamQueue + + +class WorkResult(NamedTuple): + """The value produced by a task or stream item, with new nested work.""" + + value: Any + work: Work | None = None + + +StreamItem = WorkResult +TaskResult = WorkResult + + +class WorkTask: + """A computation belonging to one or more groups. + + Compared and hashed by identity, like all graph nodes. + """ + + __slots__ = "computation", "groups" + + groups: Sequence[Group] + computation: Computation[WorkResult] + + def __init__( + self, groups: Sequence[Group], computation: Computation[WorkResult] + ) -> None: + self.groups = groups + self.computation = computation + + +class Work(NamedTuple): + """New work produced during execution.""" + + groups: Sequence[Group] = () + tasks: Sequence[WorkTask] = () + streams: Sequence[Stream] = () + + +# internal graph events + + +class _TaskSuccess(NamedTuple): + task: WorkTask + result: WorkResult + + +class _TaskFailure(NamedTuple): + task: WorkTask + error: BaseException + + +class _StreamItems(NamedTuple): + stream: Stream + items: Sequence[StreamItem] + handled: Event # pacing: the pump proceeds only after handling + + +class _StreamSuccess(NamedTuple): + stream: Stream + + +class _StreamFailure(NamedTuple): + stream: Stream + error: BaseException + + +# public work queue events + + +class GroupValuesEvent(NamedTuple): + """All values produced for a completed group.""" + + group: Group + values: Sequence[Any] + + +class GroupSuccessEvent(NamedTuple): + """A group completed successfully, promoting new work.""" + + group: Group + new_groups: Sequence[Group] + new_streams: Sequence[Stream] + + +class GroupFailureEvent(NamedTuple): + """A group failed with an error.""" + + group: Group + error: BaseException + + +class StreamValuesEvent(NamedTuple): + """New values produced by a stream, promoting new work.""" + + stream: Stream + values: Sequence[Any] + new_groups: Sequence[Group] + new_streams: Sequence[Stream] + + +class StreamSuccessEvent(NamedTuple): + """A stream completed successfully.""" + + stream: Stream + + +class StreamFailureEvent(NamedTuple): + """A stream failed with an error.""" + + stream: Stream + error: BaseException + + +class WorkQueueTerminationEvent(NamedTuple): + """All work has been delivered; the work queue terminates.""" + + +WorkQueueEvent: TypeAlias = ( + GroupValuesEvent + | GroupSuccessEvent + | GroupFailureEvent + | StreamValuesEvent + | StreamSuccessEvent + | StreamFailureEvent + | WorkQueueTerminationEvent +) + + +class _GroupNode: + """Bookkeeping for a group in the work graph.""" + + __slots__ = "child_groups", "pending", "tasks" + + child_groups: list[Group] + tasks: dict[WorkTask, None] # used as ordered set + pending: int + + def __init__(self) -> None: + self.child_groups = [] + self.tasks = {} + self.pending = 0 + + +class _TaskNode: + """Bookkeeping for a started task in the work graph.""" + + __slots__ = "child_streams", "value" + + value: Any + child_streams: list[Stream] + + def __init__(self) -> None: + self.value = _UNSET + self.child_streams = [] + + +class WorkQueue: + """Dependency-graph scheduler for incremental delivery. + + Integrates the work produced during execution into a dependency graph, + starts tasks and streams when their group becomes deliverable, prunes + empty groups, removes failed subtrees, and translates task and stream + completions into an asynchronous sequence of work queue event batches. + + For internal use only. + """ + + initial_groups: Sequence[Group] + initial_streams: Sequence[Stream] + + def __init__(self, initial_work: Work | None = None) -> None: + """Initialize the work queue with the initially produced work.""" + self._root_groups: dict[Group, None] = {} # used as ordered set + self._root_streams: dict[Stream, None] = {} # used as ordered set + self._group_nodes: dict[Group, _GroupNode] = {} + self._task_nodes: dict[WorkTask, _TaskNode] = {} + self._channel: Queue[Any] = Queue() + self._stopped = False + self._pump_tasks: set[Task[None]] = set() + + new_groups, new_streams = self._maybe_integrate_work(initial_work) + non_empty_initial_root_groups = self._prune_empty_groups(new_groups) + # Initialize root groups and streams at startup to prepare for + # cancellation prior to starting the work queue. + for group in non_empty_initial_root_groups: + self._root_groups[group] = None + for stream in new_streams: + self._root_streams[stream] = None + self.initial_groups = non_empty_initial_root_groups + self.initial_streams = new_streams + + async def events(self) -> AsyncIterator[Sequence[WorkQueueEvent]]: + """Iterate over the batches of work queue events. + + Starts the root groups and streams on the first iteration step. + """ + for group in list(self._root_groups): + self._start_group(group) + for stream in list(self._root_streams): + self._start_stream(stream) + channel = self._channel + while not self._stopped: + work_queue_events: list[WorkQueueEvent] = [] + graph_event = await channel.get() + while True: + if graph_event is not _STOP: + # Handling a graph event may synchronously push further + # graph events (e.g. when a completed group releases a + # child group whose tasks complete synchronously); these + # are picked up and handled within the same batch. + work_queue_events.extend(self._handle_graph_event(graph_event)) + try: + graph_event = channel.get_nowait() + except QueueEmpty: + break + if not self._root_groups and not self._root_streams: + self._stopped = True + work_queue_events.append(WorkQueueTerminationEvent()) + if work_queue_events: + yield work_queue_events + + async def cancel(self, reason: BaseException | None = None) -> None: + """Cancel all pending work, awaiting any asynchronous cleanup.""" + self._stopped = True + self._channel.put_nowait(_STOP) # wake up a parked event consumer + cancel_awaitables: list[Awaitable[Any]] = [] + for group in list(self._root_groups): + self._cancel_group(group, reason, cancel_awaitables) + for stream in list(self._root_streams): + self._cancel_stream(stream, reason, cancel_awaitables) + for pump_task in self._pump_tasks: + pump_task.cancel() + cancel_awaitables.extend(self._pump_tasks) + if cancel_awaitables: + await gather(*cancel_awaitables, return_exceptions=True) + + def _cancel_group( + self, + group: Group, + reason: BaseException | None, + cancel_awaitables: list[Awaitable[Any]], + ) -> None: + """Recursively cancel a group with its tasks and child groups.""" + group_node = self._group_nodes.get(group) + # defensive guard mirroring the JS version; while pending groups are + # never cancelled twice, their nodes always exist + if group_node: # pragma: no branch + for task in group_node.tasks: + self._cancel_task(task, reason, cancel_awaitables) + for child_group in group_node.child_groups: + self._cancel_group(child_group, reason, cancel_awaitables) + + def _cancel_task( + self, + task: WorkTask, + reason: BaseException | None, + cancel_awaitables: list[Awaitable[Any]], + ) -> None: + """Cancel a task with the streams produced by it.""" + abort_result = task.computation.abort(reason) + if is_awaitable(abort_result): + cancel_awaitables.append(abort_result) + task_node = self._task_nodes.get(task) + if task_node: + for child_stream in task_node.child_streams: + self._cancel_stream(child_stream, reason, cancel_awaitables) + + def _cancel_stream( + self, + stream: Stream, + reason: BaseException | None, + cancel_awaitables: list[Awaitable[Any]], + ) -> None: + """Cancel a stream.""" + abort_result = stream.queue.abort(reason) + if is_awaitable(abort_result): + cancel_awaitables.append(abort_result) + + def _push(self, graph_event: Any) -> None: + """Push a graph event to the channel unless already stopped.""" + if not self._stopped: + self._channel.put_nowait(graph_event) + + def _maybe_integrate_work( + self, work: Work | None, parent_task: WorkTask | None = None + ) -> tuple[list[Group], Sequence[Stream]]: + """Integrate new work into the graph, returning new root work.""" + if not work: + return [], [] + groups, tasks, streams = work + new_groups = self._add_groups(groups, parent_task) if groups else [] + for task in tasks: + self._add_task(task) + new_streams = self._add_streams(streams, parent_task) if streams else [] + return new_groups, new_streams + + def _add_groups( + self, original_groups: Sequence[Group], parent_task: WorkTask | None = None + ) -> list[Group]: + """Add new groups to the graph, returning the new root groups.""" + group_set = set(original_groups) + visited: set[Group] = set() + new_root_groups: list[Group] = [] + for group in original_groups: + self._add_group(group, group_set, new_root_groups, visited, parent_task) + return new_root_groups + + def _add_group( + self, + group: Group, + group_set: set[Group], + new_root_groups: list[Group], + visited: set[Group], + parent_task: WorkTask | None = None, + ) -> None: + """Add a new group to the graph, parents first.""" + if group in visited: + return + visited.add(group) + parent = group.parent + if parent is not None and parent in group_set: + self._add_group(parent, group_set, new_root_groups, visited, parent_task) + + self._group_nodes[group] = _GroupNode() + + if parent_task is None and not parent: + new_root_groups.append(group) + elif parent: + parent_node = self._group_nodes.get(parent) + if parent_node: + parent_node.child_groups.append(group) + + def _add_task(self, task: WorkTask) -> None: + """Add a new task to the graph, starting it if its group is a root.""" + for group in task.groups: + group_node = self._group_nodes.get(group) + if group_node: + group_node.tasks[task] = None + group_node.pending += 1 + if group in self._root_groups: + self._start_task(task) + + def _add_streams( + self, streams: Sequence[Stream], parent_task: WorkTask | None = None + ) -> Sequence[Stream]: + """Add new streams, attaching them to the producing task if any.""" + if not parent_task: + return streams + task_node = self._task_nodes.get(parent_task) + if task_node: + task_node.child_streams.extend(streams) + return [] + + def _prune_empty_groups( + self, + new_groups: Sequence[Group], + non_empty_new_groups: list[Group] | None = None, + ) -> list[Group]: + """Prune empty groups, promoting their non-empty descendants.""" + if non_empty_new_groups is None: + non_empty_new_groups = [] + group_nodes = self._group_nodes + for new_group in new_groups: + new_group_node = group_nodes.get(new_group) + if new_group_node: + if new_group_node.pending: + non_empty_new_groups.append(new_group) + else: + del group_nodes[new_group] + self._prune_empty_groups( + new_group_node.child_groups, non_empty_new_groups + ) + return non_empty_new_groups + + def _start_new_work( + self, new_groups: Sequence[Group], new_streams: Sequence[Stream] + ) -> None: + """Promote new groups and streams to root and start them.""" + for group in new_groups: + self._root_groups[group] = None + self._start_group(group) + for stream in new_streams: + self._root_streams[stream] = None + self._start_stream(stream) + + def _start_group(self, group: Group) -> None: + """Start all tasks of a group.""" + group_node = self._group_nodes.get(group) + # defensive guard mirroring the JS version; groups are only + # promoted right after pruning, when their nodes always exist + if group_node: # pragma: no branch + for task in list(group_node.tasks): + self._start_task(task) + + def _start_task(self, task: WorkTask) -> None: + """Start a task, pushing a graph event when its computation settles.""" + task_nodes = self._task_nodes + if task in task_nodes: + return + task_nodes[task] = _TaskNode() + try: + result = task.computation.result() + except Exception as error: + self._push(_TaskFailure(task, error)) + return + if is_awaitable(result): + future: Future[WorkResult] = ensure_future(result) + + def settle_task(future: Future[WorkResult]) -> None: + if future.cancelled(): + self._push(_TaskFailure(task, CancelledError())) + return + error = future.exception() + if error is None: + self._push(_TaskSuccess(task, future.result())) + else: + self._push(_TaskFailure(task, error)) + + future.add_done_callback(settle_task) + else: + self._push(_TaskSuccess(task, cast("WorkResult", result))) + + def _start_stream(self, stream: Stream) -> None: + """Start pumping the batches of a stream into the channel.""" + + async def pump() -> None: + try: + async for items in stream.queue.batches(): + handled = Event() + self._push(_StreamItems(stream, list(items), handled)) + # Wait until the items were handled before proceeding, so + # that work spawned by them is scheduled before the next + # batch and before the success of this stream. + await handled.wait() + except Exception as error: + self._push(_StreamFailure(stream, error)) + else: + self._push(_StreamSuccess(stream)) + + pump_task = ensure_future(pump()) + pump_tasks = self._pump_tasks + pump_tasks.add(pump_task) + pump_task.add_done_callback(pump_tasks.discard) + + def _handle_graph_event(self, graph_event: Any) -> Sequence[WorkQueueEvent]: + """Translate a single graph event into work queue events.""" + if isinstance(graph_event, _TaskSuccess): + return self._task_success(graph_event) + if isinstance(graph_event, _TaskFailure): + return self._task_failure(graph_event) + if isinstance(graph_event, _StreamItems): + return self._stream_items(graph_event) + root_streams = self._root_streams + if isinstance(graph_event, _StreamSuccess): + stream = graph_event.stream + # check whether already delivered within _stream_items() + if stream in root_streams: + del root_streams[stream] + return [StreamSuccessEvent(stream)] + return [] + # _StreamFailure + stream, error = graph_event + root_streams.pop(stream, None) + return [StreamFailureEvent(stream, error)] + + def _task_success( + self, graph_event: _TaskSuccess + ) -> Sequence[GroupValuesEvent | GroupSuccessEvent]: + """Handle a task success, finishing groups that become complete.""" + task, result = graph_event + value, work = result + task_node = self._task_nodes.get(task) + if task_node: + task_node.value = value + self._maybe_integrate_work(work, task) + + group_events: list[GroupValuesEvent | GroupSuccessEvent] = [] + new_groups: list[Group] = [] + new_streams: list[Stream] = [] + root_groups = self._root_groups + for group in task.groups: + group_node = self._group_nodes.get(group) + if group_node: + group_node.pending -= 1 + if group in root_groups and not group_node.pending: + ( + group_values_event, + group_success_event, + child_new_groups, + child_new_streams, + ) = self._finish_group_success(group, group_node) + if group_values_event: + group_events.append(group_values_event) + group_events.append(group_success_event) + new_groups.extend(child_new_groups) + new_streams.extend(child_new_streams) + + self._start_new_work(new_groups, new_streams) + return group_events + + def _task_failure(self, graph_event: _TaskFailure) -> Sequence[GroupFailureEvent]: + """Handle a task failure, removing all groups it belongs to.""" + task, error = graph_event + self._task_nodes.pop(task, None) + group_failure_events: list[GroupFailureEvent] = [] + for group in task.groups: + group_node = self._group_nodes.get(group) + if group_node: + group_failure_events.append( + self._finish_group_failure(group, group_node, error) + ) + return group_failure_events + + def _stream_items( + self, graph_event: _StreamItems + ) -> Sequence[StreamValuesEvent | StreamSuccessEvent]: + """Handle new stream items, integrating the work they produced.""" + stream, items, handled = graph_event + values: list[Any] = [] + new_groups: list[Group] = [] + new_streams: list[Stream] = [] + for value, work in items: + item_new_groups, item_new_streams = self._maybe_integrate_work(work) + non_empty_new_groups = self._prune_empty_groups(item_new_groups) + self._start_new_work(non_empty_new_groups, item_new_streams) + values.append(value) + new_groups.extend(non_empty_new_groups) + new_streams.extend(item_new_streams) + handled.set() # resume the pump of this stream + stream_values_event = StreamValuesEvent(stream, values, new_groups, new_streams) + + # queues allow peeking ahead to see if the stream has stopped + if stream.queue.is_stopped(): + self._root_streams.pop(stream, None) + return [stream_values_event, StreamSuccessEvent(stream)] + return [stream_values_event] + + def _finish_group_success( + self, group: Group, group_node: _GroupNode + ) -> tuple[ + GroupValuesEvent | None, + GroupSuccessEvent, + Sequence[Group], + Sequence[Stream], + ]: + """Finish a successfully completed group, promoting its children.""" + del self._group_nodes[group] + values: list[Any] = [] + new_streams: list[Stream] = [] + task_nodes = self._task_nodes + for task in list(group_node.tasks): + task_node = task_nodes.get(task) + if task_node: # pragma: no branch + value = task_node.value + if value is not _UNSET: # pragma: no branch + values.append(value) + new_streams.extend(task_node.child_streams) + self._remove_task(task) + new_groups = self._prune_empty_groups(group_node.child_groups) + del self._root_groups[group] + return ( + GroupValuesEvent(group, values) if values else None, + GroupSuccessEvent(group, new_groups, new_streams), + new_groups, + new_streams, + ) + + def _finish_group_failure( + self, group: Group, group_node: _GroupNode, error: BaseException + ) -> GroupFailureEvent: + """Finish a failed group, removing its whole subtree.""" + self._remove_group(group, group_node) + self._root_groups.pop(group, None) + return GroupFailureEvent(group, error) + + def _remove_group(self, group: Group, group_node: _GroupNode) -> None: + """Remove a group with its tasks and child groups from the graph.""" + group_nodes = self._group_nodes + del group_nodes[group] + for task in list(group_node.tasks): + if all(task_group not in group_nodes for task_group in task.groups): + self._remove_task(task) + for child_group in group_node.child_groups: + child_group_node = group_nodes.get(child_group) + if child_group_node: + self._remove_group(child_group, child_group_node) + + def _remove_task(self, task: WorkTask) -> None: + """Remove a task from all its groups and from the graph.""" + group_nodes = self._group_nodes + for group in task.groups: + group_node = group_nodes.get(group) + if group_node: + group_node.tasks.pop(task, None) + self._task_nodes.pop(task, None) diff --git a/src/graphql/execution/incremental_graph.py b/src/graphql/execution/incremental_graph.py deleted file mode 100644 index 40375b42..00000000 --- a/src/graphql/execution/incremental_graph.py +++ /dev/null @@ -1,357 +0,0 @@ -"""Incremental Graphs.""" - -from __future__ import annotations - -from asyncio import ( - Future, - Task, - ensure_future, - gather, - get_running_loop, - isfuture, - sleep, -) -from collections import deque -from contextlib import suppress -from typing import TYPE_CHECKING, Any, cast - -from ..pyutils import BoxedAwaitableOrValue, Undefined, is_awaitable -from .types import ( - StreamItemsRecordResult, - StreamItemsResult, - StreamRecord, - is_deferred_fragment_record, - is_pending_execution_group, -) - -if TYPE_CHECKING: - from collections.abc import Awaitable, Generator, Iterable, Sequence - - from ..error.graphql_error import GraphQLError - from .types import ( - DeferredFragmentRecord, - DeliveryGroup, - IncrementalDataRecord, - IncrementalDataRecordResult, - PendingExecutionGroup, - SuccessfulExecutionGroup, - ) - -__all__ = ["IncrementalGraph"] - - -class IncrementalGraph: - """Helper class to execute incremental Graphs. - - For internal use only. - """ - - _root_nodes: dict[DeliveryGroup, None] - _completed_queue: list[IncrementalDataRecordResult] - _next_queue: list[Future[Iterable[IncrementalDataRecordResult] | None]] - - _tasks: set[Task[Any]] - - def __init__(self) -> None: - """Initialize the IncrementalGraph.""" - self._root_nodes = {} - self._completed_queue = [] - self._next_queue = [] - self._tasks = set() - - def get_new_root_nodes( - self, incremental_data_records: Sequence[IncrementalDataRecord] - ) -> list[DeliveryGroup]: - """Get new root nodes.""" - initial_result_children: dict[DeliveryGroup, None] = {} - self._add_incremental_data_records( - incremental_data_records, None, initial_result_children - ) - return self._promote_non_empty_to_root(initial_result_children) - - def add_completed_successful_execution_group( - self, successful_execution_group: SuccessfulExecutionGroup - ) -> None: - """Add a completed successful execution group.""" - pending_group = successful_execution_group.pending_execution_group - incremental_records = successful_execution_group.incremental_data_records - deferred_records = pending_group.deferred_fragment_records - for deferred_record in deferred_records: - pending_groups = deferred_record.pending_execution_groups - successful_groups = deferred_record.successful_execution_groups - del pending_groups[pending_group] - successful_groups[successful_execution_group] = None - - if incremental_records is not None: - self._add_incremental_data_records(incremental_records, deferred_records) - - def current_completed_batch( - self, - ) -> Generator[IncrementalDataRecordResult, None, None]: - """Yield the current completed batch of incremental data record results.""" - queue = self._completed_queue - while queue: - yield queue.pop(0) - if not self._root_nodes: - self.abort() - - def next_completed_batch( - self, - ) -> Future[Iterable[IncrementalDataRecordResult] | None]: - """Return a future that resolves to the next completed batch.""" - loop = get_running_loop() - future: Future[Iterable[IncrementalDataRecordResult] | None] = ( - loop.create_future() - ) - self._next_queue.append(future) - return future - - def abort(self) -> None: - """Abort the incremental graph execution.""" - for resolve in self._next_queue: - if not resolve.cancelled(): # pragma: no cover - resolve.set_result(None) - - def has_next(self) -> bool: - """Check if there are more results to process.""" - return bool(self._root_nodes) - - def complete_deferred_fragment( - self, - deferred_fragment_record: DeferredFragmentRecord, - ) -> ( - tuple[ - list[DeliveryGroup], - list[SuccessfulExecutionGroup], - ] - | None - ): - """Complete a deferred fragment.""" - if ( - deferred_fragment_record not in self._root_nodes - or deferred_fragment_record.pending_execution_groups - ): - return None - successful_execution_groups = list( - deferred_fragment_record.successful_execution_groups - ) - del self._root_nodes[deferred_fragment_record] - for successful_execution_group in successful_execution_groups: - pending_execution_group = successful_execution_group.pending_execution_group - deferred_records = pending_execution_group.deferred_fragment_records - for other_deferred_fragment_record in deferred_records: - with suppress(KeyError): - del other_deferred_fragment_record.successful_execution_groups[ - successful_execution_group - ] - new_root_nodes = self._promote_non_empty_to_root( - deferred_fragment_record.children - ) - return new_root_nodes, successful_execution_groups - - def remove_deferred_fragment( - self, - deferred_fragment_record: DeferredFragmentRecord, - ) -> bool: - """Check if deferred fragment exists and remove it in that case.""" - if deferred_fragment_record not in self._root_nodes: - return False - del self._root_nodes[deferred_fragment_record] - return True - - def remove_stream(self, stream_record: StreamRecord) -> None: - """Remove a stream record as no longer pending.""" - del self._root_nodes[stream_record] - - async def stop_incremental_data(self) -> None: - """Stop the delivery and execution of incremental data. - - Cancels all pending requests for the next completed batch and all still - running incremental execution tasks, waiting until their cancellation - has settled. - """ - for future in self._next_queue: - future.cancel() - tasks = list(self._tasks) - if tasks: - for task in tasks: - task.cancel() - await gather(*tasks, return_exceptions=True) - - def _add_incremental_data_records( - self, - incremental_data_records: Sequence[IncrementalDataRecord], - parents: Sequence[DeferredFragmentRecord] | None = None, - initial_result_children: dict[DeliveryGroup, None] | None = None, - ) -> None: - """Add incremental data records.""" - for incremental_data_record in incremental_data_records: - if is_pending_execution_group(incremental_data_record): - deferred_records = incremental_data_record.deferred_fragment_records - for deferred_fragment_record in deferred_records: - self._add_deferred_fragment_node( - deferred_fragment_record, initial_result_children - ) - deferred_fragment_record.pending_execution_groups[ - incremental_data_record - ] = None - if self._completes_root_node(incremental_data_record): - self._on_execution_group(incremental_data_record) - elif parents is None: - if initial_result_children is None: # pragma: no cover - msg = "Invalid state while adding incremental data records." - raise RuntimeError(msg) - initial_result_children[ - cast("StreamRecord", incremental_data_record) - ] = None - else: - for parent in parents: - self._add_deferred_fragment_node(parent, initial_result_children) - parent.children[cast("StreamRecord", incremental_data_record)] = ( - None - ) - - def _promote_non_empty_to_root( - self, maybe_empty_new_root_nodes: dict[DeliveryGroup, None] - ) -> list[DeliveryGroup]: - """Promote non-empty nodes to root nodes.""" - new_root_nodes: list[DeliveryGroup] = [] - # use a deque to simulate how JavaScripts iterates over a changing set - unprocessed_nodes = deque(maybe_empty_new_root_nodes) - while unprocessed_nodes: - node = unprocessed_nodes.popleft() - if is_deferred_fragment_record(node): - pending_execution_groups = node.pending_execution_groups - if pending_execution_groups: - for pending_execution_group in pending_execution_groups: - if not self._completes_root_node(pending_execution_group): - self._on_execution_group(pending_execution_group) - self._root_nodes[node] = None - new_root_nodes.append(node) - continue - for child in node.children: - if child not in maybe_empty_new_root_nodes: # pragma: no branch - maybe_empty_new_root_nodes[cast("StreamRecord", child)] = None - unprocessed_nodes.append(child) - else: - self._root_nodes[node] = None - new_root_nodes.append(cast("StreamRecord", node)) - self._add_task(self._on_stream_items(cast("StreamRecord", node))) - return new_root_nodes - - def _completes_root_node( - self, pending_execution_group: PendingExecutionGroup - ) -> bool: - """Check whether the given record completes a root node.""" - root_nodes = self._root_nodes - deferred_records = pending_execution_group.deferred_fragment_records - return any(record in root_nodes for record in deferred_records) - - def _add_deferred_fragment_node( - self, - deferred_fragment_record: DeferredFragmentRecord, - initial_result_children: dict[DeliveryGroup, None] | None = None, - ) -> None: - """Add a deferred fragment node.""" - if deferred_fragment_record in self._root_nodes: - return - parent = deferred_fragment_record.parent - if parent is None: - if initial_result_children is None: - # The fragment has already been removed from the root nodes - # (e.g. its execution group failed on a resolver error) while a - # sibling execution group that references its subtree is still - # completing. There is nothing to attach the record to and its - # subtree is no longer being delivered, so drop it rather than - # treating this as an invalid state. Mirrors the - # `future.cancelled()` guard in `_enqueue` — a race with a - # stopping/removed consumer. - return - initial_result_children[deferred_fragment_record] = None - return - parent.children[deferred_fragment_record] = None - self._add_deferred_fragment_node(parent, initial_result_children) - - def _on_execution_group( - self, pending_execution_group: PendingExecutionGroup - ) -> None: - """Handle deferred grouped field set record.""" - completed_execution_group = pending_execution_group.result - if not isinstance(completed_execution_group, BoxedAwaitableOrValue): - completed_execution_group = completed_execution_group() - value = completed_execution_group.value - if is_awaitable(value): - - async def await_and_enqueue() -> None: - self._enqueue(await value) - - self._add_task(await_and_enqueue()) - else: - self._enqueue(value) - - async def _on_stream_items(self, stream_record: StreamRecord) -> None: - """Handle stream items.""" - enqueue = self._enqueue - items: list[Any] = [] - errors: list[GraphQLError] = [] - incremental_data_records: list[IncrementalDataRecord] = [] - stream_item_queue = stream_record.stream_item_queue - while True: - try: - stream_item_record = stream_item_queue.pop(0) - except IndexError: # pragma: no cover - break - result = ( - stream_item_record.value - if isinstance(stream_item_record, BoxedAwaitableOrValue) - else stream_item_record().value - ) - if isfuture(result): - if items: - enqueue( - StreamItemsResult( - stream_record, - incremental_data_records, - StreamItemsRecordResult(items, errors or None), - ) - ) - items = [] - errors = [] - incremental_data_records = [] - await sleep(0) # allow other tasks to run - result = await result - if result.item is Undefined: - if items: - enqueue( - StreamItemsResult( - stream_record, - incremental_data_records, - StreamItemsRecordResult(items, errors or None), - ) - ) - enqueue(StreamItemsResult(stream_record, errors=result.errors or None)) - return - items.append(result.item) - if result.errors: - errors.extend(result.errors) - if result.incremental_data_records: - incremental_data_records.extend(result.incremental_data_records) - - def _enqueue(self, completed: IncrementalDataRecordResult) -> None: - """Enqueue completed incremental data record result.""" - self._completed_queue.append(completed) - next_queue = self._next_queue - while next_queue: - future = next_queue.pop(0) - if future.cancelled(): # pragma: no cover - # defensive guard against a race with a stopping consumer - continue - future.set_result(self.current_completed_batch()) - break - - def _add_task(self, awaitable: Awaitable[Any]) -> None: - """Add the given task to the tasks set for later execution.""" - tasks = self._tasks - task = ensure_future(awaitable) - tasks.add(task) - task.add_done_callback(tasks.discard) diff --git a/src/graphql/execution/incremental_publisher.py b/src/graphql/execution/incremental_publisher.py deleted file mode 100644 index 7466aa50..00000000 --- a/src/graphql/execution/incremental_publisher.py +++ /dev/null @@ -1,358 +0,0 @@ -"""Incremental Publisher""" - -from __future__ import annotations - -from asyncio import FIRST_COMPLETED, ensure_future, sleep, wait -from contextlib import suppress -from typing import ( - TYPE_CHECKING, - Any, - NamedTuple, - Protocol, - cast, -) - -from ..pyutils import is_awaitable -from .incremental_graph import IncrementalGraph -from .types import ( - CompletedResult, - ExperimentalIncrementalExecutionResults, - IncrementalDeferResult, - IncrementalStreamResult, - InitialIncrementalExecutionResult, - PendingResult, - SubsequentIncrementalExecutionResult, - is_cancellable_stream_record, - is_completed_execution_group, - is_failed_execution_group, -) - -if TYPE_CHECKING: - from collections.abc import AsyncGenerator, Iterable, Sequence - - from ..error import GraphQLError - from ..pyutils import AbortSignal - from .types import ( - CancellableStreamRecord, - CompletedExecutionGroup, - DeferredFragmentRecord, - DeliveryGroup, - IncrementalDataRecord, - IncrementalDataRecordResult, - IncrementalResult, - StreamItemsResult, - SuccessfulExecutionGroup, - ) - -__all__ = [ - "IncrementalPublisher", - "IncrementalPublisherContext", - "build_incremental_response", -] - -suppress_key_error = suppress(KeyError) - - -class IncrementalPublisherContext(Protocol): - """The context for incremental publishing.""" - - abort_signal: AbortSignal | None - cancellable_streams: set[CancellableStreamRecord] | None - - def abort_error(self) -> Exception: - """Return the exception to raise when execution has been aborted.""" - ... # pragma: no cover - - async def cancel_incremental_work(self) -> None: - """Cancel all pending incremental work and close the stream sources.""" - ... # pragma: no cover - - def run_async_work_finished_hook(self) -> None: - """Run the hook signaling that all asynchronous work has finished.""" - ... # pragma: no cover - - -class SubsequentIncrementalExecutionResultContext(NamedTuple): - """The context for subsequent incremental execution results.""" - - pending: list[PendingResult] - incremental: list[IncrementalResult] - completed: list[CompletedResult] - - -class IncrementalPublisher: - """Publish incremental results. - - This class is used to publish incremental results to the client, enabling - semi-concurrent execution while preserving result order. - - For internal use only. - """ - - _context: IncrementalPublisherContext - _next_id: int - _incremental_graph: IncrementalGraph - - def __init__(self, context: IncrementalPublisherContext) -> None: - self._context = context - self._next_id = 0 - self._incremental_graph = IncrementalGraph() - - def build_response( - self, - data: dict[str, Any], - errors: list[GraphQLError] | None, - incremental_data_records: Sequence[IncrementalDataRecord], - ) -> ExperimentalIncrementalExecutionResults: - """Build response.""" - new_root_nodes = self._incremental_graph.get_new_root_nodes( - incremental_data_records - ) - - pending = self._to_pending_results(new_root_nodes) - - initial_result = InitialIncrementalExecutionResult( - data, - errors or None, - pending=pending, - has_next=True, - ) - - return ExperimentalIncrementalExecutionResults( - initial_result, self._subscribe() - ) - - def _to_pending_results( - self, new_root_nodes: list[DeliveryGroup] - ) -> list[PendingResult]: - """Convert pending sources to pending results.""" - pending_results: list[PendingResult] = [] - for node in new_root_nodes: - id_ = self._get_next_id() - node.id = id_ - path = node.path - pending_results.append( - PendingResult(id_, path.as_list() if path else [], node.label) - ) - return pending_results - - def _get_next_id(self) -> str: - """Get the next ID for pending results.""" - id_ = self._next_id - self._next_id += 1 - return str(id_) - - async def _subscribe( - self, - ) -> AsyncGenerator[SubsequentIncrementalExecutionResult, None]: - """Subscribe to the incremental results.""" - incremental_graph = self._incremental_graph - check_has_next = incremental_graph.has_next - handle_completed_incremental_data = self._handle_completed_incremental_data - abort_signal = self._context.abort_signal - - try: - while True: - if abort_signal is not None and abort_signal.aborted: - raise self._context.abort_error() - - batch: Iterable[IncrementalDataRecordResult] | None = ( - incremental_graph.current_completed_batch() - ) - - while batch is not None: # pragma: no branch - context = SubsequentIncrementalExecutionResultContext([], [], []) - for completed_result in batch: - await handle_completed_incremental_data( - completed_result, context - ) - - if context.incremental or context.completed: # pragma: no branch - has_next = check_has_next() - - yield SubsequentIncrementalExecutionResult( - has_next=has_next, - pending=context.pending or None, - incremental=context.incremental or None, - completed=context.completed or None, - ) - - if not has_next: - return - - next_batch = incremental_graph.next_completed_batch() - if abort_signal is None: - batch = await next_batch - else: - # reject the pending request when the operation is aborted - abort = ensure_future(abort_signal.wait()) - try: - await wait({next_batch, abort}, return_when=FIRST_COMPLETED) - finally: - if not abort.done(): - abort.cancel() - if abort_signal.aborted: - next_batch.cancel() - raise self._context.abort_error() - batch = next_batch.result() - finally: - await self._stop_async_iterators() - self._context.run_async_work_finished_hook() - - async def _stop_async_iterators(self) -> None: - """Stop the incremental execution and finish all async iterators.""" - await self._incremental_graph.stop_incremental_data() - await self._context.cancel_incremental_work() - - async def _handle_completed_incremental_data( - self, - completed_incremental_data: IncrementalDataRecordResult, - context: SubsequentIncrementalExecutionResultContext, - ) -> None: - """Handle completed incremental data.""" - if is_completed_execution_group(completed_incremental_data): - self._handle_completed_execution_group(completed_incremental_data, context) - else: - completed_incremental_data = cast( - "StreamItemsResult", completed_incremental_data - ) - await self._handle_completed_stream_items( - completed_incremental_data, context - ) - - def _handle_completed_execution_group( - self, - completed_execution_group: CompletedExecutionGroup, - context: SubsequentIncrementalExecutionResultContext, - ) -> None: - """Handle completed deferred grouped field set result.""" - append_completed = context.completed.append - append_incremental = context.incremental.append - pending_group = completed_execution_group.pending_execution_group - if is_failed_execution_group(completed_execution_group): - remove_deferred = self._incremental_graph.remove_deferred_fragment - for deferred_fragment_record in pending_group.deferred_fragment_records: - if not remove_deferred(deferred_fragment_record): - # multiple deferred grouped field sets could error for a fragment - continue - id_ = deferred_fragment_record.id - if id_ is None: # pragma: no cover - msg = "Missing deferred fragment record identifier." - raise RuntimeError(msg) - append_completed(CompletedResult(id_, completed_execution_group.errors)) - return - - completed_execution_group = cast( - "SuccessfulExecutionGroup", completed_execution_group - ) - self._incremental_graph.add_completed_successful_execution_group( - completed_execution_group - ) - - complete_deferred = self._incremental_graph.complete_deferred_fragment - pending = context.pending - for deferred_fragment_record in pending_group.deferred_fragment_records: - completion = complete_deferred(deferred_fragment_record) - if completion is None: - continue - id_ = deferred_fragment_record.id - if id_ is None: # pragma: no cover - msg = "Missing deferred fragment record identifier." - raise RuntimeError(msg) - new_root_nodes, successful_execution_groups = completion - pending.extend(self._to_pending_results(new_root_nodes)) - for successful_execution_group in successful_execution_groups: - best_id, sub_path = self._get_best_id_and_sub_path( - id_, deferred_fragment_record, successful_execution_group - ) - result = successful_execution_group.result - incremental_entry = IncrementalDeferResult( - data=result.data, - id=best_id, - sub_path=sub_path, - errors=result.errors, - ) - append_incremental(incremental_entry) - append_completed(CompletedResult(id_)) - - async def _handle_completed_stream_items( - self, - stream_items_result: StreamItemsResult, - context: SubsequentIncrementalExecutionResultContext, - ) -> None: - """Handle completed stream.""" - stream_record = stream_items_result.stream_record - id_ = stream_record.id - if id_ is None: # pragma: no cover - msg = "Missing stream record identifier." - raise RuntimeError(msg) - incremental_graph = self._incremental_graph - if stream_items_result.errors is not None: - context.completed.append(CompletedResult(id_, stream_items_result.errors)) - incremental_graph.remove_stream(stream_record) - if is_cancellable_stream_record(stream_record): - cancellable_streams = self._context.cancellable_streams - if cancellable_streams: # pragma: no branch - cancellable_streams.discard(stream_record) - with suppress(Exception): - early_return = stream_record.early_return() - if is_awaitable(early_return): # pragma: no branch - await early_return - elif stream_items_result.result is None: - context.completed.append(CompletedResult(id_)) - incremental_graph.remove_stream(stream_record) - if is_cancellable_stream_record(stream_record): - cancellable_streams = self._context.cancellable_streams - if cancellable_streams: # pragma: no branch - cancellable_streams.discard(stream_record) - else: - result = stream_items_result.result - incremental_entry = IncrementalStreamResult( - items=result.items, id=id_, errors=result.errors - ) - context.incremental.append(incremental_entry) - incremental_data_records = stream_items_result.incremental_data_records - if incremental_data_records is not None: # pragma: no branch - new_root_nodes = incremental_graph.get_new_root_nodes( - incremental_data_records - ) - context.pending.extend(self._to_pending_results(new_root_nodes)) - await sleep(0) # allow other tasks to run - - def _get_best_id_and_sub_path( - self, - initial_id: str, - initial_deferred_fragment_record: DeferredFragmentRecord, - completed_execution_group: CompletedExecutionGroup, - ) -> tuple[str, list[str | int] | None]: - """Get the best ID and sub path for the deferred grouped field set result.""" - path = initial_deferred_fragment_record.path - max_length = len(path.as_list()) if path else 0 - best_id = initial_id - pending_group = completed_execution_group.pending_execution_group - for deferred_fragment_record in pending_group.deferred_fragment_records: - if deferred_fragment_record is initial_deferred_fragment_record: - continue - id_ = deferred_fragment_record.id - if id_ is None: - continue # pragma: no cover - fragment_path = deferred_fragment_record.path - length = len(fragment_path.as_list()) if fragment_path else 0 - if length > max_length: - max_length = length - best_id = id_ - sub_path = completed_execution_group.path[max_length:] or None - return best_id, sub_path - - -def build_incremental_response( - context: IncrementalPublisherContext, - result: dict[str, Any], - errors: list[GraphQLError] | None, - incremental_data_records: Sequence[IncrementalDataRecord], -) -> ExperimentalIncrementalExecutionResults: - """Build an incremental response.""" - incremental_publisher = IncrementalPublisher(context) - return incremental_publisher.build_response( - result, errors, incremental_data_records - ) diff --git a/src/graphql/execution/types.py b/src/graphql/execution/types.py index 6f1446d5..964e7a8e 100644 --- a/src/graphql/execution/types.py +++ b/src/graphql/execution/types.py @@ -2,23 +2,18 @@ from __future__ import annotations -from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator from typing import ( TYPE_CHECKING, Any, NamedTuple, TypeAlias, TypedDict, - TypeVar, ) -from ..pyutils import BoxedAwaitableOrValue, Undefined - if TYPE_CHECKING: - from typing import TypeGuard + from collections.abc import AsyncGenerator, Iterator from ..error import GraphQLError, GraphQLFormattedError - from ..pyutils import Path try: from typing import NotRequired @@ -26,11 +21,10 @@ from typing_extensions import NotRequired __all__ = [ - "DeferredFragmentRecord", - "DeliveryGroup", - "ExecutionGroupResult", + "CompletedResult", "ExecutionResult", "ExperimentalIncrementalExecutionResults", + "FormattedCompletedResult", "FormattedExecutionResult", "FormattedIncrementalDeferResult", "FormattedIncrementalResult", @@ -38,23 +32,12 @@ "FormattedInitialIncrementalExecutionResult", "FormattedPendingResult", "FormattedSubsequentIncrementalExecutionResult", - "IncrementalDataRecord", "IncrementalDeferResult", "IncrementalResult", "IncrementalStreamResult", "InitialIncrementalExecutionResult", "PendingResult", - "StreamItemRecord", - "StreamItemResult", - "StreamItemsRecordResult", - "StreamItemsResult", "SubsequentIncrementalExecutionResult", - "SuccessfulExecutionGroup", - "is_cancellable_stream_record", - "is_completed_execution_group", - "is_deferred_fragment_record", - "is_failed_execution_group", - "is_pending_execution_group", ] @@ -341,29 +324,16 @@ class FormattedSubsequentIncrementalExecutionResult(TypedDict): extensions: NotRequired[dict[str, Any]] -class ExecutionGroupResult: - """Execution group result.""" - - errors: list[GraphQLError] | None - data: dict[str, Any] - - __slots__ = "data", "errors" - - def __init__( - self, data: dict[str, Any], errors: list[GraphQLError] | None = None - ) -> None: - self.data = data - self.errors = errors - - -class IncrementalDeferResult(ExecutionGroupResult): # noqa: PLW1641 +class IncrementalDeferResult: # noqa: PLW1641 """Incremental deferred execution result""" + data: dict[str, Any] id: str sub_path: list[str | int] | None + errors: list[GraphQLError] | None extensions: dict[str, Any] | None - __slots__ = "extensions", "id", "sub_path" + __slots__ = "data", "errors", "extensions", "id", "sub_path" def __init__( self, @@ -447,31 +417,16 @@ class FormattedInitialIncrementalExecutionResult(TypedDict): extensions: NotRequired[dict[str, Any]] -class StreamItemsRecordResult: - """Stream items record result.""" - - errors: list[GraphQLError] | None - items: list[Any] - - __slots__ = "errors", "items" - - def __init__( - self, - items: list[Any], - errors: list[GraphQLError] | None = None, - ) -> None: - self.items = items - self.errors = errors - - -class IncrementalStreamResult(StreamItemsRecordResult): +class IncrementalStreamResult: """Incremental streamed execution result""" + items: list[Any] id: str sub_path: list[str | int] | None + errors: list[GraphQLError] | None extensions: dict[str, Any] | None - __slots__ = "extensions", "id", "sub_path" + __slots__ = "errors", "extensions", "id", "items", "sub_path" def __init__( self, @@ -568,8 +523,6 @@ class FormattedIncrementalStreamResult(TypedDict): extensions: NotRequired[dict[str, Any]] -T = TypeVar("T") # declare T for generic aliases - IncrementalResult: TypeAlias = IncrementalDeferResult | IncrementalStreamResult FormattedIncrementalResult: TypeAlias = ( @@ -687,241 +640,3 @@ class FormattedCompletedResult(TypedDict): id: str errors: NotRequired[list[GraphQLFormattedError]] - - -def is_pending_execution_group( - incremental_data_record: IncrementalDataRecord, -) -> TypeGuard[PendingExecutionGroup]: - """Check if the incremental data record is a pending execution group.""" - return isinstance(incremental_data_record, PendingExecutionGroup) - - -class SuccessfulExecutionGroup: - """Successful execution group""" - - pending_execution_group: PendingExecutionGroup - path: list[str | int] - result: ExecutionGroupResult - incremental_data_records: list[IncrementalDataRecord] | None - errors: None = None - - __slots__ = ( - "incremental_data_records", - "path", - "pending_execution_group", - "result", - ) - - def __init__( - self, - pending_execution_group: PendingExecutionGroup, - path: list[str | int], - result: ExecutionGroupResult, - incremental_data_records: list[IncrementalDataRecord] | None = None, - ) -> None: - self.pending_execution_group = pending_execution_group - self.path = path - self.result = result - self.incremental_data_records = incremental_data_records - - -class FailedExecutionGroup: - """Failed execution group""" - - pending_execution_group: PendingExecutionGroup - path: list[str | int] - errors: list[GraphQLError] - result: None = None - - __slots__ = "errors", "path", "pending_execution_group" - - def __init__( - self, - pending_execution_group: PendingExecutionGroup, - path: list[str | int], - errors: list[GraphQLError], - ) -> None: - self.pending_execution_group = pending_execution_group - self.path = path - self.errors = errors - - -def is_failed_execution_group( - completed_execution_group: CompletedExecutionGroup, -) -> TypeGuard[FailedExecutionGroup]: - """Check if the completed execution group is a failed execution group.""" - return isinstance(completed_execution_group, FailedExecutionGroup) - - -CompletedExecutionGroup: TypeAlias = SuccessfulExecutionGroup | FailedExecutionGroup - - -def is_completed_execution_group( - incremental_data_record_result: IncrementalDataRecordResult, -) -> TypeGuard[CompletedExecutionGroup]: - """Check if the subsequent result is a deferred grouped field set result.""" - return isinstance( - incremental_data_record_result, - SuccessfulExecutionGroup | FailedExecutionGroup, - ) - - -ThunkIncrementalResult: TypeAlias = ( - BoxedAwaitableOrValue[T] | Callable[[], BoxedAwaitableOrValue[T]] -) - - -class PendingExecutionGroup: - """Pending execution group""" - - deferred_fragment_records: list[DeferredFragmentRecord] - result: ThunkIncrementalResult[CompletedExecutionGroup] - - __slots__ = "deferred_fragment_records", "result" - - def __init__( - self, - deferred_fragment_records: list[DeferredFragmentRecord], - result: ThunkIncrementalResult[CompletedExecutionGroup], - ) -> None: - self.result = result - self.deferred_fragment_records = deferred_fragment_records - - -class DeferredFragmentRecord: - """Deferred fragment record""" - - path: Path | None - label: str | None - id: str | None - parent: DeferredFragmentRecord | None - pending_execution_groups: dict[PendingExecutionGroup, None] - successful_execution_groups: dict[SuccessfulExecutionGroup, None] - children: dict[DeliveryGroup, None] - - __slots__ = ( - "children", - "id", - "label", - "parent", - "path", - "pending_execution_groups", - "successful_execution_groups", - ) - - def __init__( - self, - path: Path | None = None, - label: str | None = None, - parent: DeferredFragmentRecord | None = None, - ) -> None: - self.path = path - self.label = label - self.parent = parent - self.id = None - self.pending_execution_groups = {} - self.successful_execution_groups = {} - self.children = {} - - def __repr__(self) -> str: - name = self.__class__.__name__ - args: list[str] = [] - if self.path: - args.append(f"path={self.path.as_list()!r}") - if self.label: - args.append(f"label={self.label!r}") - if self.parent: - args.append("parent") - return f"{name}({', '.join(args)})" - - -def is_deferred_fragment_record( - delivery_group: DeliveryGroup, -) -> TypeGuard[DeferredFragmentRecord]: - """Check if the delivery group is a deferred fragment record.""" - return isinstance(delivery_group, DeferredFragmentRecord) - - -class StreamItemResult(NamedTuple): - """Stream item result""" - - item: Any = Undefined - incremental_data_records: list[IncrementalDataRecord] | None = None - errors: list[GraphQLError] | None = None - - -StreamItemRecord: TypeAlias = ThunkIncrementalResult[StreamItemResult] - - -class StreamRecord: - """Stream record""" - - stream_item_queue: list[StreamItemRecord] - path: Path - label: str | None - id: str | None - - __slots__ = "id", "label", "path", "stream_item_queue" - - def __init__( - self, - stream_item_queue: list[StreamItemRecord], - path: Path, - label: str | None = None, - ) -> None: - self.stream_item_queue = stream_item_queue - self.path = path - self.label = label - self.id = None - - def __repr__(self) -> str: - name = self.__class__.__name__ - args: list[str] = [ - f"stream_item_queue[{len(self.stream_item_queue)}]", - f"path={self.path.as_list()!r}", - ] - if self.label: - args.append(f"label={self.label!r}") - return f"{name}({', '.join(args)})" - - -DeliveryGroup: TypeAlias = DeferredFragmentRecord | StreamRecord - - -class CancellableStreamRecord(StreamRecord): - """Cancellable stream record""" - - early_return: Callable[[], Awaitable[None]] - - __slots__ = ("early_return",) - - def __init__( - self, - early_return: Callable[[], Awaitable[None]], - stream_item_queue: list[StreamItemRecord], - path: Path, - label: str | None = None, - ) -> None: - super().__init__(stream_item_queue, path, label) - self.early_return = early_return - - -def is_cancellable_stream_record( - delivery_group: DeliveryGroup, -) -> TypeGuard[CancellableStreamRecord]: - """Check if the delivery group is a cancellable stream record.""" - return isinstance(delivery_group, CancellableStreamRecord) - - -class StreamItemsResult(NamedTuple): - """Stream items result""" - - stream_record: StreamRecord - incremental_data_records: list[IncrementalDataRecord] | None = None - result: StreamItemsRecordResult | None = None - errors: list[GraphQLError] | None = None - - -IncrementalDataRecord: TypeAlias = PendingExecutionGroup | StreamRecord - -IncrementalDataRecordResult: TypeAlias = CompletedExecutionGroup | StreamItemsResult diff --git a/src/graphql/pyutils/__init__.py b/src/graphql/pyutils/__init__.py index c4658871..2d6c8766 100644 --- a/src/graphql/pyutils/__init__.py +++ b/src/graphql/pyutils/__init__.py @@ -10,7 +10,6 @@ from .abort_signal import AbortController, AbortError, AbortSignal from .async_reduce import async_reduce -from .boxed_awaitable_or_value import BoxedAwaitableOrValue from .gather_with_cancel import gather_with_cancel from .convert_case import camel_to_snake, snake_to_camel from .cached_property import cached_property @@ -44,7 +43,6 @@ "AbortError", "AbortSignal", "AwaitableOrValue", - "BoxedAwaitableOrValue", "Description", "FrozenError", "Path", diff --git a/src/graphql/pyutils/boxed_awaitable_or_value.py b/src/graphql/pyutils/boxed_awaitable_or_value.py deleted file mode 100644 index 8b92e135..00000000 --- a/src/graphql/pyutils/boxed_awaitable_or_value.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Boxed Awaitable or Value""" - -from __future__ import annotations - -from asyncio import CancelledError, Future, ensure_future, isfuture -from contextlib import suppress -from typing import TYPE_CHECKING, Generic, TypeVar - -if TYPE_CHECKING: - from collections.abc import Awaitable - -__all__ = ["BoxedAwaitableOrValue"] - -T = TypeVar("T") - - -class BoxedAwaitableOrValue(Generic[T]): - """Container for an Awaitable or a Value that updates itself. - - A BoxedAwaitableOrValue is a container for a value or Awaitable where the value - will be updated when the Awaitable has been awaited. - """ - - __slots__ = "_future", "_value" - - _value: T | Future[T] - - def __init__(self, value: T | Awaitable[T]) -> None: - """Initialize the BoxedAwaitableOrValue with the given value or Awaitable.""" - try: - value = ensure_future(value) # type: ignore - except TypeError: - pass - else: - value.add_done_callback(self._update_value) - self._value = value # type: ignore - - @property - def value(self) -> T: - """Get the current value.""" - value = self._value - if isfuture(value) and value.done(): - self._value = value = value.result() - return value # type: ignore - - @property - def pending_future(self) -> Future[T] | None: - """Get the still pending Future, or None if the value is already settled.""" - value = self._value - if isfuture(value) and not value.done(): - return value - return None - - def _update_value(self, value: Future[T]) -> None: - """Update the boxed value when the Awaitable is done.""" - with suppress(CancelledError): - self._value = value.result() diff --git a/tests/execution/incremental/__init__.py b/tests/execution/incremental/__init__.py new file mode 100644 index 00000000..897a7e24 --- /dev/null +++ b/tests/execution/incremental/__init__.py @@ -0,0 +1 @@ +"""Tests for graphql.execution.incremental""" diff --git a/tests/execution/incremental/test_computation.py b/tests/execution/incremental/test_computation.py new file mode 100644 index 00000000..c344f27f --- /dev/null +++ b/tests/execution/incremental/test_computation.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +from asyncio import CancelledError, Event, Future, sleep +from typing import TYPE_CHECKING, Any + +import pytest + +from graphql.execution.incremental import Computation + +if TYPE_CHECKING: + from collections.abc import Callable + +pytestmark = pytest.mark.anyio + + +class Spy: + """Wrap a function, counting its calls.""" + + def __init__(self, fn: Callable[..., Any]) -> None: + self.fn = fn + self.call_count = 0 + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self.call_count += 1 + return self.fn(*args, **kwargs) + + +def describe_computation(): + def can_return_a_result(): + computation = Computation(lambda: {"value": 123}) + + assert computation.result() == {"value": 123} + + def can_be_started_manually(): + computation = Computation(lambda: {"value": 123}) + + computation.prime() + assert computation.result() == {"value": 123} + + def only_runs_once_when_started_multiple_times(): + run_spy = Spy(lambda: {"value": "done"}) + computation = Computation(run_spy) + + computation.prime() + computation.prime() + computation.prime() + results = [computation.result() for _ in range(3)] + + assert results == [{"value": "done"}] * 3 + assert run_spy.call_count == 1 + + async def stores_async_result_via_result(): + async def run() -> dict[str, str]: + await sleep(0) + return {"value": "done"} + + run_spy = Spy(run) + computation = Computation(run_spy) + + computation.prime() + computation.prime() + computation.prime() + assert await computation.result() == {"value": "done"} + results = [computation.result() for _ in range(3)] + + assert results == [{"value": "done"}] * 3 + assert run_spy.call_count == 1 + + def stores_sync_error_in_result(): + def run() -> None: + raise RuntimeError("failure") + + run_spy = Spy(run) + computation = Computation(run_spy) + + computation.prime() # does not raise + with pytest.raises(RuntimeError, match="failure"): + computation.result() + with pytest.raises(RuntimeError, match="failure"): + computation.result() + assert run_spy.call_count == 1 + + async def stores_async_error_in_result(): + async def run() -> None: + await sleep(0) + raise RuntimeError("failure") + + run_spy = Spy(run) + computation = Computation(run_spy) + + computation.prime() # does not raise + with pytest.raises(RuntimeError, match="failure"): + await computation.result() + with pytest.raises(RuntimeError, match="failure"): + computation.result() + assert run_spy.call_count == 1 + + def can_be_aborted_before_running(): + on_abort_spy = Spy(lambda _reason: None) + run_spy = Spy(lambda: {"value": 123}) + computation = Computation(run_spy, on_abort_spy) + + computation.abort(RuntimeError("Cancelled!")) + with pytest.raises(RuntimeError, match="Cancelled!"): + computation.result() + assert on_abort_spy.call_count == 0 + assert run_spy.call_count == 0 + + def cannot_be_aborted_after_running_synchronously(): + on_abort_spy = Spy(lambda _reason: None) + computation = Computation(lambda: {"value": 123}, on_abort_spy) + + computation.prime() + computation.abort() + assert computation.result() == {"value": 123} + assert on_abort_spy.call_count == 0 + + def cannot_be_aborted_after_erroring_synchronously(): + def run() -> None: + raise RuntimeError("failure") + + on_abort_spy = Spy(lambda _reason: None) + computation: Computation[Any] = Computation(run, on_abort_spy) + + computation.prime() + computation.abort() + with pytest.raises(RuntimeError, match="failure"): + computation.result() + assert on_abort_spy.call_count == 0 + + async def can_be_aborted_while_running_asynchronously(): + async def run() -> None: + # may be cancelled before it even starts running + await Event().wait() # pragma: no cover + + on_abort_spy = Spy(lambda _reason: None) + computation: Computation[Any] = Computation(run, on_abort_spy) + + computation.prime() + computation.abort(RuntimeError("Cancelled!")) + assert on_abort_spy.call_count == 1 + with pytest.raises(RuntimeError, match="Cancelled!"): + computation.result() + await sleep(0) # let the cancelled future settle + + async def returns_async_abort_cleanup_while_running(): + async def run() -> None: + # may be cancelled before it even starts running + await Event().wait() # pragma: no cover + + cleanup_future: Future[None] = Future() + computation: Computation[Any] = Computation(run, lambda _reason: cleanup_future) + + computation.prime() + abort_result = computation.abort() + assert abort_result is cleanup_future + assert not cleanup_future.done() + cleanup_future.set_result(None) + await cleanup_future + await sleep(0) # let the cancelled future settle + + def can_be_aborted_with_a_provided_reason_before_running(): + abort_reason = RuntimeError("aborted") + computation = Computation(lambda: {"value": 123}) + + computation.abort(abort_reason) + with pytest.raises(RuntimeError, match="aborted") as exc_info: + computation.result() + assert exc_info.value is abort_reason + + async def forwards_abort_reason_to_on_abort_while_running_asynchronously(): + async def run() -> None: + # may be cancelled before it even starts running + await Event().wait() # pragma: no cover + + abort_reason = RuntimeError("aborted") + recorded_reasons: list[BaseException | None] = [] + computation: Computation[Any] = Computation(run, recorded_reasons.append) + + computation.prime() + computation.abort(abort_reason) + assert recorded_reasons == [abort_reason] + with pytest.raises(RuntimeError, match="aborted"): + computation.result() + await sleep(0) # let the cancelled future settle + + def aborting_without_reason_raises_a_cancelled_error(): + computation = Computation(lambda: {"value": 123}) + + computation.abort() + with pytest.raises(CancelledError): + computation.result() + + def repeated_aborts_are_ignored(): + computation = Computation(lambda: {"value": 123}) + + computation.abort(RuntimeError("first")) + computation.abort(RuntimeError("second")) + with pytest.raises(RuntimeError, match="first"): + computation.result() + + async def can_be_aborted_while_running_without_abort_callback(): + async def run() -> None: + # may be cancelled before it even starts running + await Event().wait() # pragma: no cover + + computation: Computation[Any] = Computation(run) + + computation.prime() + assert computation.abort(RuntimeError("Cancelled!")) is None + with pytest.raises(RuntimeError, match="Cancelled!"): + computation.result() + await sleep(0) # let the cancelled future settle + + async def cannot_be_aborted_after_completing_asynchronously(): + async def run() -> dict[str, int]: + return {"value": 123} + + on_abort_spy = Spy(lambda _reason: None) + computation: Computation[Any] = Computation(run, on_abort_spy) + + computation.prime() + assert await computation.result() == {"value": 123} + computation.abort() + assert computation.result() == {"value": 123} + assert on_abort_spy.call_count == 0 + + async def rejects_when_the_pending_future_is_cancelled_externally(): + async def run() -> None: + # may be cancelled before it even starts running + await Event().wait() # pragma: no cover + + computation: Computation[Any] = Computation(run) + + future = computation.result() + assert isinstance(future, Future) + future.cancel() + with pytest.raises(CancelledError): + await future + with pytest.raises(CancelledError): + computation.result() + + async def abort_wins_when_racing_async_completion(): + async def run() -> dict[str, int]: + return {"value": 123} + + computation: Computation[Any] = Computation(run) + + computation.prime() + await sleep(0) # future completes, but memoization is still queued + computation.abort(RuntimeError("Cancelled!")) + await sleep(0) # run the queued memoization callback + with pytest.raises(RuntimeError, match="Cancelled!"): + computation.result() diff --git a/tests/execution/test_defer.py b/tests/execution/incremental/test_defer.py similarity index 82% rename from tests/execution/test_defer.py rename to tests/execution/incremental/test_defer.py index 0ddb46c0..f1479fa0 100644 --- a/tests/execution/test_defer.py +++ b/tests/execution/incremental/test_defer.py @@ -9,7 +9,6 @@ from graphql.execution import ( AbortedGraphQLExecutionError, CompletedResult, - DeferredFragmentRecord, ExecutionResult, ExperimentalIncrementalExecutionResults, IncrementalDeferResult, @@ -20,6 +19,7 @@ execute, experimental_execute_incrementally, ) +from graphql.execution.incremental import DeliveryGroup from graphql.language import DocumentNode, parse from graphql.pyutils import AbortController, AbortError, Path, is_awaitable from graphql.type import ( @@ -109,6 +109,7 @@ class Friend(NamedTuple): { "b": GraphQLField(b), "someField": GraphQLField(GraphQLString), + "nonNullErrorField": GraphQLField(GraphQLNonNull(GraphQLString)), }, ) @@ -138,6 +139,29 @@ class Friend(NamedTuple): schema = GraphQLSchema(query) +user_type = GraphQLObjectType("User", {"id": GraphQLField(GraphQLID)}) + +todo_type = GraphQLObjectType( + "Todo", + { + "id": GraphQLField(GraphQLID), + "items": GraphQLField(GraphQLList(GraphQLString)), + "author": GraphQLField(user_type), + }, +) + +cancellation_schema = GraphQLSchema( + GraphQLObjectType( + "Query", + { + "todo": GraphQLField(todo_type), + "blocker": GraphQLField(GraphQLString), + "scalarList": GraphQLField(GraphQLList(GraphQLString)), + "slowScalarList": GraphQLField(GraphQLList(GraphQLString)), + }, + ) +) + class Resolvers: """Various resolver functions for testing""" @@ -525,14 +549,12 @@ def can_compare_subsequent_incremental_execution_result(): "completed": completed, } - def can_print_deferred_fragment_record(): - """Can print a DeferredFragmentRecord""" - record = DeferredFragmentRecord() - assert str(record) == "DeferredFragmentRecord()" - record = DeferredFragmentRecord(Path(None, "bar", "Bar"), "foo", record) - assert ( - str(record) == "DeferredFragmentRecord(path=['bar'], label='foo', parent)" - ) + def can_print_delivery_group(): + """Can print a DeliveryGroup""" + group = DeliveryGroup(None, None, None) + assert str(group) == "DeliveryGroup()" + group = DeliveryGroup(Path(None, "bar", "Bar"), "foo", group) + assert str(group) == "DeliveryGroup(path=['bar'], label='foo', parent)" @pytest.mark.parametrize("early_execution", [False, True]) async def can_defer_fragments_containing_scalar_types(early_execution): @@ -1950,17 +1972,17 @@ async def nulls_cross_defer_boundaries_null_first(): assert result == [ { "data": {"a": {}}, - "pending": [{"id": "0", "path": []}, {"id": "1", "path": ["a"]}], + "pending": [{"id": "0", "path": ["a"]}, {"id": "1", "path": []}], "hasNext": True, }, { "incremental": [ - {"data": {"b": {"c": {}}}, "id": "1"}, - {"data": {"d": "d"}, "id": "1", "subPath": ["b", "c"]}, + {"data": {"b": {"c": {}}}, "id": "0"}, + {"data": {"d": "d"}, "id": "0", "subPath": ["b", "c"]}, ], "completed": [ { - "id": "0", + "id": "1", "errors": [ { "message": "Cannot return null" @@ -1970,7 +1992,7 @@ async def nulls_cross_defer_boundaries_null_first(): }, ], }, - {"id": "1"}, + {"id": "0"}, ], "hasNext": False, }, @@ -2164,18 +2186,18 @@ async def nulls_cross_defer_boundaries_value_first(): assert result == [ { "data": {"a": {}}, - "pending": [{"id": "0", "path": []}, {"id": "1", "path": ["a"]}], + "pending": [{"id": "0", "path": ["a"]}, {"id": "1", "path": []}], "hasNext": True, }, { "incremental": [ - {"data": {"b": {"c": {}}}, "id": "1"}, - {"data": {"d": "d"}, "id": "0", "subPath": ["a", "b", "c"]}, + {"data": {"b": {"c": {}}}, "id": "0"}, + {"data": {"d": "d"}, "id": "1", "subPath": ["a", "b", "c"]}, ], "completed": [ - {"id": "0"}, + {"id": "1"}, { - "id": "1", + "id": "0", "errors": [ { "message": "Cannot return null" @@ -2190,6 +2212,144 @@ async def nulls_cross_defer_boundaries_value_first(): }, ] + async def nulls_cross_defer_boundaries_failed_fragment_with_slower_shared_groups(): + """Nulls cross defer boundaries, failed fragment with slower shared children + + Nulls cross defer boundaries, failed fragment with slower shared child + execution groups. + """ + document = parse( + """ + query { + ... @defer { + a { + someField + nonNullErrorField + b { + c { + d + } + } + } + } + a { + ... @defer { + someField + b { + e { + f + } + } + } + } + } + """ + ) + + async def some_field(_info): + return "someField" + + result = await complete( + document, + {"a": {"b": {"c": {"d": "d"}, "e": {"f": "f"}}, "someField": some_field}}, + ) + assert result == [ + { + "data": {"a": {}}, + "pending": [{"id": "0", "path": ["a"]}, {"id": "1", "path": []}], + "hasNext": True, + }, + { + "completed": [ + { + "id": "1", + "errors": [ + { + "message": "Cannot return null" + " for non-nullable field a.nonNullErrorField.", + "locations": [{"line": 6, "column": 19}], + "path": ["a", "nonNullErrorField"], + }, + ], + }, + ], + "hasNext": True, + }, + { + "incremental": [ + {"data": {"b": {}, "someField": "someField"}, "id": "0"}, + {"data": {"e": {"f": "f"}}, "id": "0", "subPath": ["b"]}, + ], + "completed": [{"id": "0"}], + "hasNext": False, + }, + ] + + async def handles_cancelling_child_deferred_fragments_if_parent_fragment_fails(): + """Handles cancelling child deferred fragments if parent fragment fails""" + document = parse( + """ + query { + ... @defer { + a { + someField + b { + c { + nonNullErrorField + } + } + } + ... @defer { + a { + someField + } + } + } + a { + ... @defer { + b { + c { + d + } + } + } + } + } + """ + ) + result = await complete( + document, + {"a": {"b": {"c": {"d": "d"}}, "someField": "someField"}}, + ) + assert result == [ + { + "data": {"a": {}}, + "pending": [{"id": "0", "path": ["a"]}, {"id": "1", "path": []}], + "hasNext": True, + }, + { + "incremental": [ + {"data": {"b": {"c": {}}}, "id": "0"}, + {"data": {"d": "d"}, "id": "0", "subPath": ["b", "c"]}, + ], + "completed": [ + { + "id": "1", + "errors": [ + { + "message": "Cannot return null" + " for non-nullable field c.nonNullErrorField.", + "locations": [{"line": 8, "column": 23}], + "path": ["a", "b", "c", "nonNullErrorField"], + }, + ], + }, + {"id": "0"}, + ], + "hasNext": False, + }, + ] + async def filters_a_payload_with_a_null_that_cannot_be_merged(): """Filters a payload with a null that cannot be merged""" document = parse( @@ -2232,21 +2392,21 @@ async def filters_a_payload_with_a_null_that_cannot_be_merged(): assert result == [ { "data": {"a": {}}, - "pending": [{"id": "0", "path": []}, {"id": "1", "path": ["a"]}], + "pending": [{"id": "0", "path": ["a"]}, {"id": "1", "path": []}], "hasNext": True, }, { "incremental": [ - {"data": {"b": {"c": {}}}, "id": "1"}, - {"data": {"d": "d"}, "id": "1", "subPath": ["b", "c"]}, + {"data": {"b": {"c": {}}}, "id": "0"}, + {"data": {"d": "d"}, "id": "0", "subPath": ["b", "c"]}, ], - "completed": [{"id": "1"}], + "completed": [{"id": "0"}], "hasNext": True, }, { "completed": [ { - "id": "0", + "id": "1", "errors": [ { "message": "Cannot return null" @@ -2261,8 +2421,48 @@ async def filters_a_payload_with_a_null_that_cannot_be_merged(): }, ] - async def cancels_deferred_fields_when_initial_result_exhibits_null_bubbling(): - """Cancels deferred fields when initial result exhibits null bubbling""" + async def cancels_deferred_fields_when_null_bubbling_cancels_the_defer(): + """Cancels deferred fields when initial result exhibits null bubbling + + Cancels deferred fields when initial result exhibits null bubbling, + cancelling the defer. + """ + document = parse( + """ + query { + hero { + nonNullName + ... @defer { + name + } + } + } + """ + ) + result = await complete( + document, + {"hero": {**hero, "nonNullName": lambda _info: None}}, + enable_early_execution=True, + ) + + assert result == { + "data": {"hero": None}, + "errors": [ + { + "message": "Cannot return null" + " for non-nullable field Hero.nonNullName.", + "locations": [{"line": 4, "column": 17}], + "path": ["hero", "nonNullName"], + }, + ], + } + + async def cancels_deferred_fields_when_null_bubbling_cancels_new_fields(): + """Cancels deferred fields when initial result exhibits null bubbling + + Cancels deferred fields when initial result exhibits null bubbling, + cancelling new fields. + """ document = parse( """ query { @@ -2295,6 +2495,53 @@ async def cancels_deferred_fields_when_initial_result_exhibits_null_bubbling(): ], } + async def keeps_deferred_work_outside_nulled_error_paths(): + """Keeps deferred work outside nulled error paths""" + document = parse( + """ + query { + a { + ... @defer { + someField + } + nonNullErrorField + } + g { + ... @defer { + h + } + } + } + """ + ) + result = await complete( + document, + { + "a": {"someField": "someField", "nonNullErrorField": None}, + "g": {"h": "value"}, + }, + ) + assert result == [ + { + "data": {"a": None, "g": {}}, + "errors": [ + { + "message": "Cannot return null" + " for non-nullable field a.nonNullErrorField.", + "locations": [{"line": 7, "column": 17}], + "path": ["a", "nonNullErrorField"], + }, + ], + "pending": [{"id": "0", "path": ["g"]}], + "hasNext": True, + }, + { + "incremental": [{"data": {"h": "value"}, "id": "0"}], + "completed": [{"id": "0"}], + "hasNext": False, + }, + ] + async def stops_late_initial_path_completion_before_deferred_response(): """Stops late initial-path completion before publishing a deferred response""" document = parse( @@ -3043,6 +3290,165 @@ async def original_execute_function_throws_error_if_deferred_and_not_all_is_sync " multiple payloads (due to @defer or @stream directive)" ) + async def allows_deferred_execution_when_passed_abort_signal_if_not_aborted(): + """Should allow deferred execution when passed abortSignal, if not aborted""" + abort_controller = AbortController() + document = parse( + """ + query { + todo { + id + ... on Todo @defer { + author { + id + } + } + } + } + """ + ) + + result = experimental_execute_incrementally( + cancellation_schema, + document, + {"todo": {"author": {"id": "1"}}}, + abort_signal=abort_controller.signal, + ) + assert isinstance(result, ExperimentalIncrementalExecutionResults) + + assert result.initial_result == { + "data": {"todo": {"id": None}}, + "pending": [{"id": "0", "path": ["todo"]}], + "hasNext": True, + } + + iterator = result.subsequent_results + payload1 = await anext(iterator) + assert payload1 == { + "incremental": [{"data": {"author": {"id": "1"}}, "id": "0"}], + "completed": [{"id": "0"}], + "hasNext": False, + } + + with pytest.raises(StopAsyncIteration): + await anext(iterator) + + async def stops_deferred_execution_when_aborted(): + """Should stop deferred execution when aborted""" + abort_controller = AbortController() + document = parse( + """ + query { + todo { + id + ... on Todo @defer { + author { + id + } + } + } + } + """ + ) + + author_calls = 0 + + def author(_info): + nonlocal author_calls + author_calls += 1 # pragma: no cover + return {"id": "1"} # pragma: no cover + + async def todo(_info): + # never reached: the execution is aborted before the todo resolves + return {"id": "1", "author": author} # pragma: no cover + + awaitable_result = experimental_execute_incrementally( + cancellation_schema, + document, + {"todo": todo}, + abort_signal=abort_controller.signal, + ) + assert is_awaitable(awaitable_result) + result_task = ensure_future(awaitable_result) + + abort_controller.abort() + + with pytest.raises( + AbortedGraphQLExecutionError, match="This operation was aborted" + ): + await result_task + assert author_calls == 0 + + async def stops_deferred_execution_when_aborted_mid_execution(): + """Should stop deferred execution when aborted mid-execution""" + abort_controller = AbortController() + document = parse( + """ + query { + ... on Query @defer { + todo { + id + author { + id + } + } + } + } + """ + ) + + author_calls = 0 + + async def author(_info): + nonlocal author_calls + author_calls += 1 # pragma: no cover + return {"id": "1"} # pragma: no cover + + async def todo(_info): + # never reached: the execution is aborted before the todo resolves + return {"id": "1", "author": author} # pragma: no cover + + result = experimental_execute_incrementally( + cancellation_schema, + document, + {"todo": todo}, + abort_signal=abort_controller.signal, + ) + # the initial result has no non-deferred fields and completes early + assert isinstance(result, ExperimentalIncrementalExecutionResults) + + abort_controller.abort() + + with pytest.raises(AbortError, match="This operation was aborted"): + async for _patch in result.subsequent_results: + pass # pragma: no cover + assert author_calls == 0 + + async def cancels_pending_deferred_execution_groups(): + """Cancels pending deferred execution groups""" + abort_controller = AbortController() + document = parse("{ scalarList ... @defer { slowScalarList } }") + + slow_future: Future[Any] = Future() + + result = experimental_execute_incrementally( + cancellation_schema, + document, + { + "scalarList": lambda _info: ["a"], + "slowScalarList": lambda _info: slow_future, + }, + enable_early_execution=True, + abort_signal=abort_controller.signal, + ) + assert isinstance(result, ExperimentalIncrementalExecutionResults) + + iterator = result.subsequent_results + abort_controller.abort() + + with pytest.raises(AbortError, match="This operation was aborted"): + await anext(iterator) + async def cancels_pending_deferred_tasks_with_async_child_stream_cleanup(): """Cancels pending deferred tasks with async child stream cleanup @@ -3052,25 +3458,6 @@ async def cancels_pending_deferred_tasks_with_async_child_stream_cleanup(): only after the cleanup has settled, so that no cleanup coroutine is left behind unawaited. """ - user_type = GraphQLObjectType("User", {"id": GraphQLField(GraphQLID)}) - todo_type = GraphQLObjectType( - "Todo", - { - "id": GraphQLField(GraphQLID), - "items": GraphQLField(GraphQLList(GraphQLString)), - "author": GraphQLField(user_type), - }, - ) - cancellation_schema = GraphQLSchema( - GraphQLObjectType( - "Query", - { - "todo": GraphQLField(todo_type), - "blocker": GraphQLField(GraphQLString), - }, - ) - ) - abort_controller = AbortController() document = parse( """ @@ -3164,6 +3551,97 @@ async def author(_info): await anext_task assert items_source.aclose_finished + async def ignores_deferred_payloads_resolved_after_cancellation(): + """Should ignore deferred payloads resolved after cancellation""" + abort_controller = AbortController() + document = parse( + """ + query { + todo { + id + ... @defer { + author { + id + } + } + } + } + """ + ) + + author_started = Event() + author_future: Future[Any] = Future() + + def author(_info): + author_started.set() + return author_future + + result = experimental_execute_incrementally( + cancellation_schema, + document, + {"todo": {"id": "todo", "author": author}}, + abort_signal=abort_controller.signal, + enable_early_execution=True, + ) + assert isinstance(result, ExperimentalIncrementalExecutionResults) + + iterator = result.subsequent_results + next_task = ensure_future(anext(iterator)) + + await author_started.wait() + abort_controller.abort() + + author_future.set_result({"id": "author"}) # the late result is ignored + + with pytest.raises(AbortError, match="This operation was aborted"): + await next_task + + async def ignores_deferred_errors_after_cancellation(): + """Should ignore deferred errors after cancellation""" + abort_controller = AbortController() + document = parse( + """ + query { + todo { + id + ... @defer { + author { + id + } + } + } + } + """ + ) + + author_started = Event() + author_future: Future[Any] = Future() + + def author(_info): + author_started.set() + return author_future + + result = experimental_execute_incrementally( + cancellation_schema, + document, + {"todo": {"id": "todo", "author": author}}, + abort_signal=abort_controller.signal, + enable_early_execution=True, + ) + assert isinstance(result, ExperimentalIncrementalExecutionResults) + + iterator = result.subsequent_results + next_task = ensure_future(anext(iterator)) + + await author_started.wait() + abort_controller.abort() + + # the late error is ignored and absorbed instead of being reported + author_future.set_exception(RuntimeError("late error")) + + with pytest.raises(AbortError, match="This operation was aborted"): + await next_task + def describe_defer_directive_with_errors_and_nested_defer(): """Regression tests for graphql-python/graphql-core#271.""" @@ -3288,3 +3766,133 @@ async def resolve_obj(_source, _info) -> dict: (b_id,) = [id_ for id_, label in label_of_id.items() if label == "b"] fast_entries = [i for i in incremental if i.get("data") == {"fast": "fast"}] assert fast_entries == [{"data": {"fast": "fast"}, "id": b_id}] + + +def describe_defer_and_stream_on_spec_invalid_queries(): + """Regression tests for graphql-python/graphql-core#272.""" + + @pytest.mark.timeout(1) + @pytest.mark.parametrize("early_execution", [False, True]) + async def terminates_when_streamed_field_conflicts_with_deferred_selection( + early_execution, + ): + """A spec-invalid ``@defer``/``@stream`` combination must not hang. + + A streamed list field without a subselection that also merge-conflicts + with an unstreamed selection of the same field inside a sibling + ``@defer`` used to deadlock incremental delivery whenever another field + of that sibling resolved later than the streamed field: the incremental + graph never converged to a terminal state, so ``subsequent_results`` + never yielded a final payload with ``hasNext: False``. Spec-valid + queries were not affected, but ``experimental_execute_incrementally`` + may be called without a prior validation pass and must terminate on + invalid input, too. See issue #272. + """ + + async def resolve_fast(_source, _info) -> str: + # `fast` must resolve (at least) two event loop iterations later + # than `child` and `items` to trigger the deadlock + await sleep(0) + await sleep(0) + return "fast" + + async def resolve_obj(_source, _info) -> dict: + return {} + + async def resolve_list(_source, _info) -> list: + return [{}, {}] + + obj_type: GraphQLObjectType = GraphQLObjectType( + "Obj", + lambda: { + "fast": GraphQLField(GraphQLString, resolve=resolve_fast), + "child": GraphQLField(obj_type, resolve=resolve_obj), + "items": GraphQLField(GraphQLList(obj_type), resolve=resolve_list), + }, + ) + local_schema = GraphQLSchema( + GraphQLObjectType( + "Query", {"obj": GraphQLField(obj_type, resolve=resolve_obj)} + ) + ) + # The query is spec-invalid in two ways: the streamed `items` field + # lacks a subselection, and it merge-conflicts with the unstreamed + # `items` selection inside the sibling `@defer`. Neither invalid + # feature alone triggered the deadlock, and fixing either (aliasing + # the conflict away or adding a subselection) made it disappear. + document = parse( + """ + query { + obj { + ... { + child { + child { + ... @defer(label: "L4") { items @stream(initialCount: 1) } + } + ... @defer { child { fast items } } + } + } + } + } + """ + ) + + result = experimental_execute_incrementally( + local_schema, document, {}, enable_early_execution=early_execution + ) + assert is_awaitable(result) + result = await result + assert isinstance(result, ExperimentalIncrementalExecutionResults) + initial = result.initial_result.formatted + # Draining the stream must terminate (guarded by the timeout above). + subsequent = [patch.formatted async for patch in result.subsequent_results] + payloads: list[Any] = [initial, *subsequent] + + # The stream must terminate cleanly: `hasNext` is true on every + # payload except the last, which ends the stream. + assert initial["hasNext"] is True + assert subsequent[-1]["hasNext"] is False + assert all(patch["hasNext"] is True for patch in subsequent[:-1]) + + assert initial["data"] == {"obj": {"child": {"child": {}}}} + + # Aggregate announcements, deliveries and completions across the + # whole stream; their grouping into payloads is not asserted as it + # depends on resolver scheduling and `enable_early_execution`. + pending = [p for payload in payloads for p in payload.get("pending", [])] + incremental = [ + i for payload in payloads for i in payload.get("incremental", []) + ] + completed = [c for payload in payloads for c in payload.get("completed", [])] + + # Three delivery groups are announced: the anonymous `@defer`, the + # labeled nested `@defer`, and the stream over `items`. + id_of_path = {tuple(p["path"]): p["id"] for p in pending} + label_of_path = {tuple(p["path"]): p.get("label") for p in pending} + assert label_of_path == { + ("obj", "child"): None, + ("obj", "child", "child"): "L4", + ("obj", "child", "child", "items"): None, + } + + # Every announced group completes exactly once and without errors. + assert sorted(c["id"] for c in completed) == sorted(id_of_path.values()) + assert all("errors" not in c for c in completed) + + # All three groups deliver: the initial streamed item under `L4` and + # the second item via the stream (both as empty objects, since the + # invalid selection has no subfields), and `fast` exactly once under + # the anonymous `@defer` (the conflicting unstreamed `items` selection + # merges into the streamed one instead of being delivered again). + expected_deliveries = [ + {"data": {"items": [{}]}, "id": id_of_path["obj", "child", "child"]}, + {"items": [{}], "id": id_of_path["obj", "child", "child", "items"]}, + { + "data": {"fast": "fast"}, + "id": id_of_path["obj", "child"], + "subPath": ["child"], + }, + ] + assert len(incremental) == len(expected_deliveries) + for entry in expected_deliveries: + assert entry in incremental diff --git a/tests/execution/incremental/test_incremental_executor.py b/tests/execution/incremental/test_incremental_executor.py new file mode 100644 index 00000000..02093eb2 --- /dev/null +++ b/tests/execution/incremental/test_incremental_executor.py @@ -0,0 +1,238 @@ +"""Unit tests for the incremental executor. + +These tests cover asyncio-specific seams of the incremental executor that +are not deterministically reachable via the defer/stream test suites, +especially the asynchronous cleanup paths. +""" + +from __future__ import annotations + +from asyncio import Event, sleep +from typing import TYPE_CHECKING, Any, cast + +import pytest + +from graphql import ( + GraphQLDeferDirective, + GraphQLField, + GraphQLObjectType, + GraphQLSchema, + GraphQLString, + parse, + specified_directives, +) +from graphql.execution import experimental_execute_incrementally +from graphql.execution.incremental import ( + Computation, + ExecutionGroup, + IncrementalExecutor, + ItemStream, + StreamItemQueue, + WorkResult, +) +from graphql.pyutils import RefSet, is_awaitable + +if TYPE_CHECKING: + from graphql.pyutils import Path + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +obj_type = GraphQLObjectType("Obj", {"echo": GraphQLField(GraphQLString)}) + +schema = GraphQLSchema( + GraphQLObjectType("Query", {"echo": GraphQLField(GraphQLString)}), + subscription=GraphQLObjectType( + "Subscription", + { + "echo": GraphQLField(GraphQLString), + "obj": GraphQLField(obj_type), + }, + ), + directives=[*specified_directives, GraphQLDeferDirective], +) + + +def build_executor(source: str) -> IncrementalExecutor: + executor = IncrementalExecutor.build(schema, parse(source)) + assert isinstance(executor, IncrementalExecutor) + return executor + + +def parked_stream_queue() -> StreamItemQueue: + async def produce(_queue: StreamItemQueue) -> None: + await Event().wait() # park forever + + return StreamItemQueue(produce, eager=True) + + +def parked_execution_group(executor: IncrementalExecutor) -> ExecutionGroup: + async def parked_work() -> WorkResult: + await Event().wait() # park forever + return WorkResult(None) # pragma: no cover + + async def cleanup(_reason: BaseException | None) -> None: + await sleep(0) + + computation: Computation[WorkResult] = Computation(parked_work, cleanup) + executor.prime_now(computation) + return ExecutionGroup([], computation, None) + + +def describe_incremental_executor(): + def rejects_defer_on_subscription_root_fields(): + result = experimental_execute_incrementally( + schema, parse("subscription { ... @defer { echo } }") + ) + assert result == ( + None, + [ + { + "message": "`@defer` directive not supported on subscription" + " operations. Disable `@defer` by setting the `if` argument" + " to `false`.", + } + ], + ) + + def rejects_defer_on_subscription_subfields(): + result = experimental_execute_incrementally( + schema, + parse("subscription { obj { ... @defer { echo } } }"), + {"obj": {"echo": "hello"}}, + ) + assert result == ( + {"obj": None}, + [ + { + "message": "`@defer` directive not supported on subscription" + " operations. Disable `@defer` by setting the `if` argument" + " to `false`.", + "locations": [(1, 16)], + "path": ["obj"], + } + ], + ) + + async def aborting_incremental_work_awaits_asynchronous_cleanup(): + executor = build_executor("{ echo }") + group = parked_execution_group(executor) + executor.tasks.append(group) + stream_queue = parked_stream_queue() + await sleep(0) # let the stream producer start + executor.streams.append(ItemStream(cast("Path", None), None, stream_queue, 0)) + + abort_result = executor.abort() + assert is_awaitable(abort_result) + await abort_result + + async def aborting_in_background_settles_asynchronous_cleanup(): + executor = build_executor("{ echo }") + executor.tasks.append(parked_execution_group(executor)) + executor.abort_in_background() + assert executor.background_futures + await executor.cancel_incremental_work() + + async def cancelling_incremental_work_awaits_abort_cleanup(): + executor = build_executor("{ echo }") + executor.tasks.append(parked_execution_group(executor)) + await executor.cancel_incremental_work() + + def memoizes_sub_execution_plans_per_defer_usage_set(): + executor = build_executor("{ echo }") + defer_usage_set_1: RefSet = RefSet() + defer_usage_set_2: RefSet = RefSet() + sub_executor_1 = executor.create_sub_executor(defer_usage_set_1) + sub_executor_2 = executor.create_sub_executor(defer_usage_set_2) + grouped_field_set: dict = {} + plan_1 = sub_executor_1.build_sub_execution_plan(grouped_field_set) + # planned again under a different defer usage set + plan_2 = sub_executor_2.build_sub_execution_plan(grouped_field_set) + assert plan_2 is not plan_1 + # memoized for the same defer usage set + assert sub_executor_1.build_sub_execution_plan(grouped_field_set) is plan_1 + assert sub_executor_2.build_sub_execution_plan(grouped_field_set) is plan_2 + + async def failing_execution_group_awaits_asynchronous_abort_cleanup(): + executor = build_executor("{ echo }") + sub_executor = executor.create_sub_executor() + sub_executor.tasks.append(parked_execution_group(sub_executor)) + + async def failing_fields() -> dict: + raise RuntimeError("execution group failed") + + sub_executor.execute_fields = ( # type: ignore[method-assign] + lambda *_args: failing_fields() + ) + result = sub_executor.execute_execution_group( + [], + cast("Any", schema.query_type), + None, + None, + {}, + cast("Any", None), + ) + assert is_awaitable(result) + with pytest.raises(RuntimeError, match="execution group failed"): + await result + + async def failing_awaitable_stream_item_awaits_asynchronous_abort_cleanup(): + executor = build_executor("{ echo }") + sub_executor = executor.create_sub_executor() + sub_executor.tasks.append(parked_execution_group(sub_executor)) + + async def failing_value() -> None: + raise RuntimeError("stream item failed") + + async def failing_item() -> None: + # never awaited, disposed via close() below + raise RuntimeError("stream item failed") # pragma: no cover + + sub_executor.complete_awaitable_value = ( # type: ignore[method-assign] + lambda *_args: failing_value() + ) + item = failing_item() + result = sub_executor.complete_stream_item( + cast("Path", None), + item, + [], + cast("Any", None), + cast("Any", None), + ) + assert is_awaitable(result) + with pytest.raises(RuntimeError, match="stream item failed"): + await result + item.close() # dispose of the item that was never awaited + + async def failing_stream_item_completion_awaits_asynchronous_abort_cleanup(): + executor = build_executor("{ echo }") + sub_executor = executor.create_sub_executor() + sub_executor.tasks.append(parked_execution_group(sub_executor)) + + async def failing_completion() -> None: + raise RuntimeError("completion failed") + + def raise_error(raw_error: Exception, *_args: object) -> None: + raise raw_error + + sub_executor.complete_value = ( # type: ignore[method-assign] + lambda *_args: failing_completion() + ) + sub_executor.handle_field_error = cast( # type: ignore[method-assign] + "Any", raise_error + ) + result = sub_executor.complete_stream_item( + cast("Path", None), + "item", + [], + cast("Any", None), + cast("Any", None), + ) + assert is_awaitable(result) + with pytest.raises(RuntimeError, match="completion failed"): + await result diff --git a/tests/execution/incremental/test_incremental_publisher.py b/tests/execution/incremental/test_incremental_publisher.py new file mode 100644 index 00000000..4044d308 --- /dev/null +++ b/tests/execution/incremental/test_incremental_publisher.py @@ -0,0 +1,116 @@ +"""Unit tests for the incremental publisher. + +These tests cover defensive seams of the incremental publisher that are not +reachable via the defer/stream test suites. +""" + +from __future__ import annotations + +from asyncio import CancelledError +from typing import TYPE_CHECKING + +import pytest + +from graphql.error import GraphQLError +from graphql.execution.incremental import IncrementalPublisher +from graphql.execution.incremental.incremental_publisher import ensure_graphql_error +from graphql.pyutils import AbortController + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Sequence + + from graphql.execution.incremental import WorkQueueEvent + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +class FakeWorkQueue: + """A fake work queue whose events end without a termination event.""" + + def __init__(self) -> None: + self.cancelled = False + + async def events(self) -> AsyncIterator[Sequence[WorkQueueEvent]]: + return + yield [] # pragma: no cover + + async def cancel(self, _reason: BaseException | None = None) -> None: + self.cancelled = True + + +class FakeContext: + """A fake incremental publisher context.""" + + def __init__(self) -> None: + self.abort_signal = None + self.cancelled = False + self.hook_run = False + + def abort_error(self) -> Exception: + return RuntimeError("aborted") # pragma: no cover + + async def cancel_incremental_work( + self, _reason: BaseException | None = None + ) -> None: + self.cancelled = True + + def run_async_work_finished_hook(self) -> None: + self.hook_run = True + + +def describe_incremental_publisher(): + async def stops_when_the_work_queue_ends_without_termination(): + work_queue = FakeWorkQueue() + context = FakeContext() + publisher = IncrementalPublisher() + results = [ + result + async for result in publisher._subscribe( # noqa: SLF001 + work_queue, # type: ignore[arg-type] + context, # type: ignore[arg-type] + ) + ] + assert results == [] + assert work_queue.cancelled + assert context.cancelled + assert context.hook_run + + async def stops_when_work_queue_ends_without_termination_with_abort_signal(): + work_queue = FakeWorkQueue() + context = FakeContext() + context.abort_signal = AbortController().signal # type: ignore[assignment] + publisher = IncrementalPublisher() + results = [ + result + async for result in publisher._subscribe( # noqa: SLF001 + work_queue, # type: ignore[arg-type] + context, # type: ignore[arg-type] + ) + ] + assert results == [] + assert work_queue.cancelled + assert context.cancelled + assert context.hook_run + + def describe_ensure_graphql_error(): + def passes_graphql_errors_through(): + error = GraphQLError("test") + assert ensure_graphql_error(error) is error + + def locates_plain_exceptions(): + error = RuntimeError("test") + graphql_error = ensure_graphql_error(error) + assert isinstance(graphql_error, GraphQLError) + assert graphql_error.message == "test" + assert graphql_error.original_error is error + + def converts_base_exceptions(): + error: BaseException = CancelledError() + graphql_error = ensure_graphql_error(error) + assert isinstance(graphql_error, GraphQLError) + assert graphql_error.message == "CancelledError()" diff --git a/tests/execution/test_stream.py b/tests/execution/incremental/test_stream.py similarity index 86% rename from tests/execution/test_stream.py rename to tests/execution/incremental/test_stream.py index 0a113599..77ae052f 100644 --- a/tests/execution/test_stream.py +++ b/tests/execution/incremental/test_stream.py @@ -12,11 +12,10 @@ ExecutionResult, ExperimentalIncrementalExecutionResults, IncrementalStreamResult, - StreamItemRecord, - StreamRecord, execute, experimental_execute_incrementally, ) +from graphql.execution.incremental import ItemStream, StreamItemQueue from graphql.language import DocumentNode, parse from graphql.pyutils import AbortController, AbortError, Path from graphql.type import ( @@ -244,15 +243,14 @@ def can_hash_incremental_stream_result(): IncrementalStreamResult(**modified_args(args, extensions={"baz": 1})) ) - def can_print_stream_record(): - """Can print a StreamRecord""" - queue: list[StreamItemRecord] = [] + def can_print_item_stream(): + """Can print an ItemStream""" path = Path(None, "bar", "Bar") - expected_args = "stream_item_queue[0], path=['bar']" - record = StreamRecord(queue, path) - assert str(record) == f"StreamRecord({expected_args})" - record = StreamRecord(queue, path, "foo") - assert str(record) == f"StreamRecord({expected_args}, label='foo')" + queue = StreamItemQueue(None) # type: ignore[arg-type] + stream = ItemStream(path, None, queue, 0) + assert str(stream) == "ItemStream(path=['bar'])" + stream = ItemStream(path, "foo", queue, 0) + assert str(stream) == "ItemStream(path=['bar'], label='foo')" @pytest.mark.parametrize("early_execution", [False, True]) async def can_stream_a_list_field(early_execution): @@ -278,6 +276,58 @@ async def can_stream_a_list_field(early_execution): }, ] + async def does_not_call_return_on_an_exhausted_sync_iterator(): + """Does not call ``return`` on an exhausted sync iterator + + The JavaScript version asserts that return() is not called on the + exhausted iterator; Python iterators have no equivalent method in the + iteration protocol, but an analogous close() method must likewise not + be called by the implementation. + """ + document = parse("{ scalarList @stream(initialCount: 1) }") + + values = ["apple", "banana", "coconut"] + + class Source: + def __init__(self): + self.index = 0 + self.close_calls = 0 + + def __iter__(self): + return self + + def __next__(self): + try: + value = values[self.index] + except IndexError: + raise StopIteration from None + finally: + self.index += 1 + return value + + def close(self): + self.close_calls += 1 # pragma: no cover + + source = Source() + + result = await complete(document, {"scalarList": source}) + assert result == [ + { + "data": {"scalarList": ["apple"]}, + "pending": [{"id": "0", "path": ["scalarList"]}], + "hasNext": True, + }, + { + "incremental": [ + {"items": ["banana", "coconut"], "id": "0"}, + ], + "completed": [{"id": "0"}], + "hasNext": False, + }, + ] + assert source.close_calls == 0 + assert source.index == 4 + async def can_use_default_value_of_initial_count(): """Can use default value of initialCount""" document = parse("{ scalarList @stream }") @@ -524,9 +574,14 @@ async def await_friend(f): }, { "incremental": [ - {"items": [{"name": "Luke", "id": "1"}], "id": "0"}, - {"items": [{"name": "Han", "id": "2"}], "id": "0"}, - {"items": [{"name": "Leia", "id": "3"}], "id": "0"}, + { + "items": [ + {"name": "Luke", "id": "1"}, + {"name": "Han", "id": "2"}, + {"name": "Leia", "id": "3"}, + ], + "id": "0", + }, ], "completed": [{"id": "0"}], "hasNext": False, @@ -764,7 +819,7 @@ async def await_friend(f, i): { "incremental": [ { - "items": [None], + "items": [None, {"name": "Leia", "id": "3"}], "id": "0", "errors": [ { @@ -774,7 +829,6 @@ async def await_friend(f, i): } ], }, - {"items": [{"name": "Leia", "id": "3"}], "id": "0"}, ], "completed": [{"id": "0"}], "hasNext": False, @@ -807,9 +861,14 @@ async def friend_list(_info): }, { "incremental": [ - {"items": [{"name": "Luke", "id": "1"}], "id": "0"}, - {"items": [{"name": "Han", "id": "2"}], "id": "0"}, - {"items": [{"name": "Leia", "id": "3"}], "id": "0"}, + { + "items": [ + {"name": "Luke", "id": "1"}, + {"name": "Han", "id": "2"}, + {"name": "Leia", "id": "3"}, + ], + "id": "0", + }, ], "completed": [{"id": "0"}], "hasNext": False, @@ -1184,6 +1243,88 @@ async def friend_list(_info): }, ] + async def drains_sync_iterators_with_later_promises_when_null_bubbles(): + """Drains sync iterators with later promises when null bubbles past the stream + + The JavaScript version asserts that there is no unhandled rejection; + the asyncio analog is that the pending awaitable collected from the + drained iterator is settled in the background, so that its late error + is absorbed instead of being reported as never retrieved. + """ + document = parse( + """ + query { + nonNullFriendList @stream(initialCount: 1) { + name + } + } + """ + ) + + rejected = False + + async def delayed_reject(): + nonlocal rejected + await sleep(0) + rejected = True + raise RuntimeError("third bad") + + values = [friends[0], None, delayed_reject()] + + class Source: + def __init__(self): + self.index = 0 + self.close_calls = 0 + + def __iter__(self): + return self + + def __next__(self): + try: + value = values[self.index] + except IndexError: + raise StopIteration from None + finally: + self.index += 1 + return value + + def close(self): + self.close_calls += 1 # pragma: no cover + + source = Source() + + result = await complete(document, {"nonNullFriendList": source}) + assert result == [ + { + "data": {"nonNullFriendList": [{"name": "Luke"}]}, + "pending": [{"id": "0", "path": ["nonNullFriendList"]}], + "hasNext": True, + }, + { + "completed": [ + { + "id": "0", + "errors": [ + { + "message": "Cannot return null for non-nullable field" + " Query.nonNullFriendList.", + "locations": [{"line": 3, "column": 15}], + "path": ["nonNullFriendList", 1], + }, + ], + }, + ], + "hasNext": False, + }, + ] + assert source.close_calls == 0 + assert source.index == 4 + + # the pending awaitable from the drained iterator settles in background + for _ in range(2): + await sleep(0) + assert rejected + async def handles_error_thrown_in_complete_value_after_initial_count_is_reached(): """Handles errors thrown by completeValue after initialCount is reached""" document = parse( @@ -1260,7 +1401,7 @@ def get_friends(_info): "hasNext": False, "incremental": [ { - "items": [None], + "items": [None, {"nonNullName": "Han"}], "id": "0", "errors": [ { @@ -1270,7 +1411,6 @@ def get_friends(_info): } ], }, - {"items": [{"nonNullName": "Han"}], "id": "0"}, ], "completed": [{"id": "0"}], }, @@ -1315,7 +1455,7 @@ def get_friends(_info): { "incremental": [ { - "items": [None], + "items": [None, {"nonNullName": "Han"}], "id": "0", "errors": [ { @@ -1325,7 +1465,6 @@ def get_friends(_info): } ], }, - {"items": [{"nonNullName": "Han"}], "id": "0"}, ], "completed": [{"id": "0"}], "hasNext": False, @@ -1414,7 +1553,10 @@ def get_friends(_info): { "incremental": [ { - "items": [None], + "items": [ + None, + {"metadata": {"value": "later"}, "nonNullName": "Han"}, + ], "id": "0", "errors": [ { @@ -1424,12 +1566,6 @@ def get_friends(_info): }, ], }, - { - "items": [ - {"metadata": {"value": "later"}, "nonNullName": "Han"} - ], - "id": "0", - }, ], "completed": [{"id": "0"}], "hasNext": False, @@ -1591,7 +1727,7 @@ async def get_friends(_info): { "incremental": [ { - "items": [None], + "items": [None, {"nonNullName": "Han"}], "id": "0", "errors": [ { @@ -1601,7 +1737,6 @@ async def get_friends(_info): } ], }, - {"items": [{"nonNullName": "Han"}], "id": "0"}, ], "completed": [{"id": "0"}], "hasNext": False, @@ -1951,6 +2086,71 @@ async def friend_list(_info): }, ] + async def cancels_async_stream_items_when_null_bubbles_past_the_stream(): + """Cancels async stream items when null bubbles past the stream + + In the asyncio implementation, the pending stream item is cancelled + when the stream is filtered out by the null bubbling, so that it can + no longer be resolved late. + """ + document = parse( + """ + query { + nestedObject { + nestedFriendList @stream(initialCount: 0) { + name + } + nonNullScalarField + } + } + """ + ) + + friends_started = Event() + non_null_future: Future[Any] = Future() + name_future: Future[str] = Future() + + def nested_friend_list(_info): + friends_started.set() + return [{"name": name_future}] + + result_future = ensure_future( + experimental_execute_incrementally( # type: ignore + schema, + document, + { + "nestedObject": { + "nestedFriendList": nested_friend_list, + "nonNullScalarField": lambda _info: non_null_future, + } + }, + enable_early_execution=True, + ) + ) + + await friends_started.wait() + await sleep(0) + non_null_future.set_result(None) + + result = await result_future + assert not isinstance(result, ExperimentalIncrementalExecutionResults) + assert result.formatted == { + "data": {"nestedObject": None}, + "errors": [ + { + "message": "Cannot return null for non-nullable field" + " NestedObject.nonNullScalarField.", + "locations": [{"line": 7, "column": 17}], + "path": ["nestedObject", "nonNullScalarField"], + } + ], + } + + # the pending stream item was cancelled instead of being resolved late + for _ in range(2): + await sleep(0) + assert name_future.cancelled() + async def filters_stream_payloads_that_are_nulled_in_a_deferred_payload(): """Filters stream payloads that are nulled in a deferred payload""" document = parse( @@ -2197,11 +2397,15 @@ async def get_friends(_info): "hasNext": True, }, { - "incremental": [{"items": [{"id": "2", "name": "Han"}], "id": "0"}], - "hasNext": True, - }, - { - "incremental": [{"items": [{"id": "3", "name": "Leia"}], "id": "0"}], + "incremental": [ + { + "items": [ + {"id": "2", "name": "Han"}, + {"id": "3", "name": "Leia"}, + ], + "id": "0", + }, + ], "completed": [{"id": "0"}], "hasNext": False, }, @@ -2209,6 +2413,28 @@ async def get_friends(_info): async def handles_overlapping_deferred_and_non_deferred_streams(): """Handles overlapping deferred and non-deferred streams""" + # NOTE: This is an *invalid* query, but it should be an *executable* query. + # Validation rejects the overlapping @stream selections. See + # https://github.com/graphql/defer-stream-wg/discussions/100. + # + # The query selects the same stream field twice. The non-deferred selection + # asks for `id`. The deferred stream selection asks for `id`, `name`, and + # an inner deferred `innerName`. + # + # Stream item completion later sees one merged stream. Without rewriting + # the stream field details to clear inherited defer usage, `name` would be + # treated as still belonging to the outer @defer instead of to the stream + # item. The bad stream payload would be `items: [{ id: '1' }]` and then + # `items: [{ id: '2' }]`; the `name` fields would be dropped even though + # the `innerName` payloads would still be emitted. + # + # `get_stream_usage` avoids that by rewriting the stream field details to + # clear inherited defer usage before those details are reused to complete + # stream items. + # + # The inner @defer is included only because it forces stream item + # completion through execution planning. Without it, a fast path executes + # all collected fields directly and this bug would not be exposed. document = parse( """ query { @@ -2222,6 +2448,9 @@ async def handles_overlapping_deferred_and_non_deferred_streams(): nestedFriendList @stream(initialCount: 0) { id name + ... @defer { + innerName: name + } } } } @@ -2250,16 +2479,111 @@ async def get_nested_friend_list(_info): "hasNext": True, }, { - "incremental": [{"items": [{"id": "1", "name": "Luke"}], "id": "0"}], + "pending": [ + {"id": "1", "path": ["nestedObject", "nestedFriendList", 0]} + ], + "incremental": [ + {"items": [{"id": "1", "name": "Luke"}], "id": "0"}, + {"data": {"innerName": "Luke"}, "id": "1"}, + ], + "completed": [{"id": "1"}], "hasNext": True, }, { - "incremental": [{"items": [{"id": "2", "name": "Han"}], "id": "0"}], - "completed": [{"id": "0"}], + "pending": [ + {"id": "2", "path": ["nestedObject", "nestedFriendList", 1]} + ], + "incremental": [ + {"items": [{"id": "2", "name": "Han"}], "id": "0"}, + {"data": {"innerName": "Han"}, "id": "2"}, + ], + "completed": [{"id": "0"}, {"id": "2"}], "hasNext": False, }, ] + async def repromotes_completed_stream_when_slower_sibling_defer_resolves_later(): + """Re-promotes a completed stream when a slower sibling defer resolves later + + The JavaScript version delivers the re-promoted stream items and the + slow deferred field as two separate payloads due to microtask timing; + the asyncio implementation delivers the same content in the same order + coalesced into a single payload. + """ + # NOTE: This is an *invalid* query, but it should be an *executable* query. + # Validation rejects the overlapping @stream selections. See + # https://github.com/graphql/defer-stream-wg/discussions/100. + document = parse( + """ + query { + nestedObject { + ... @defer { + nestedFriendList @stream { + name + } + } + ... @defer { + scalarField + nestedFriendList @stream { + name + } + } + } + } + """ + ) + + slow_field_future: Future[str] = Future() + + execute_result = experimental_execute_incrementally( + schema, + document, + { + "nestedObject": { + "nestedFriendList": lambda _info: friends, + "scalarField": lambda _info: slow_field_future, + } + }, + ) + assert isinstance(execute_result, ExperimentalIncrementalExecutionResults) + iterator = execute_result.subsequent_results + + result1 = execute_result.initial_result + assert result1 == { + "data": {"nestedObject": {}}, + "pending": [ + {"id": "0", "path": ["nestedObject"]}, + {"id": "1", "path": ["nestedObject"]}, + ], + "hasNext": True, + } + + result2 = await anext(iterator) + assert result2 == { + "pending": [{"id": "2", "path": ["nestedObject", "nestedFriendList"]}], + "incremental": [{"data": {"nestedFriendList": []}, "id": "0"}], + "completed": [{"id": "0"}], + "hasNext": True, + } + + slow_field_future.set_result("slow") + + result3 = await anext(iterator) + assert result3 == { + "incremental": [ + { + "items": [{"name": "Luke"}, {"name": "Han"}, {"name": "Leia"}], + "id": "2", + }, + {"data": {"scalarField": "slow"}, "id": "1"}, + ], + "completed": [{"id": "2"}, {"id": "1"}], + "hasNext": False, + } + + with pytest.raises(StopAsyncIteration): + await anext(iterator) + async def returns_payloads_properly_when_parent_deferred_slower_than_stream(): """Returns payloads in correct order when parent deferred slower than stream @@ -2348,7 +2672,7 @@ async def can_defer_fields_that_are_resolved_after_async_iterable_is_complete(): """ query { friendList @stream(label:"stream-label") { - ...NameFragment @defer(label: "DeferName") @defer(label: "DeferName") + ...NameFragment @defer(label: "DeferName") id } } @@ -2411,12 +2735,17 @@ async def get_friends(_info): assert result3 == { "pending": [{"id": "2", "path": ["friendList", 1], "label": "DeferName"}], "incremental": [{"items": [{"id": "2"}], "id": "0"}], - "completed": [{"id": "0"}], "hasNext": True, } result4 = await anext(iterator) assert result4 == { + "completed": [{"id": "0"}], + "hasNext": True, + } + + result5 = await anext(iterator) + assert result5 == { "incremental": [{"data": {"name": "Han"}, "id": "2"}], "completed": [{"id": "2"}], "hasNext": False, @@ -2431,7 +2760,7 @@ async def can_defer_fields_that_are_resolved_before_async_iterable_is_complete() """ query { friendList @stream(initialCount: 1, label:"stream-label") { - ...NameFragment @defer(label: "DeferName") @defer(label: "DeferName") + ...NameFragment @defer(label: "DeferName") id } } @@ -2833,6 +3162,71 @@ async def aclose(self): await iterator.aclose() assert source.aclose_calls == 0 + async def limits_stream_batches_to_the_default_capacity(): + """Limits stream batches to the default capacity (100)""" + document = parse( + """ + query { + friendList @stream { + id + } + } + """ + ) + + filled = Event() + + async def friend_list(_info): + for i in range(101): + await sleep(0) + if i == 100: # the producer parks when yielding the 101st item + filled.set() + yield friends[i % 3] + + execute_result = await experimental_execute_incrementally( # type: ignore + schema, + document, + {"friendList": friend_list}, + enable_early_execution=True, + ) + assert isinstance(execute_result, ExperimentalIncrementalExecutionResults) + iterator = execute_result.subsequent_results + + result1 = execute_result.initial_result + assert result1 == { + "data": {"friendList": []}, + "pending": [{"id": "0", "path": ["friendList"]}], + "hasNext": True, + } + + await filled.wait() # allow the producer to fill the stream queue + for _ in range(10): + await sleep(0) + + result2 = await anext(iterator) + assert result2 == { + "incremental": [ + { + "items": [{"id": str(i % 3 + 1)} for i in range(100)], + "id": "0", + } + ], + "hasNext": True, + } + + for _ in range(10): # allow the producer to push the remaining item + await sleep(0) + + result3 = await anext(iterator) + assert result3 == { + "incremental": [{"items": [{"id": "2"}], "id": "0"}], + "completed": [{"id": "0"}], + "hasNext": False, + } + + with pytest.raises(StopAsyncIteration): + await anext(iterator) + user_type = GraphQLObjectType("User", {"id": GraphQLField(GraphQLID)}) @@ -3402,6 +3796,74 @@ async def aclose(self): await sleep(0) assert author_calls == 0 + @pytest.mark.timeout(1) + async def stops_when_stream_queue_is_back_pressured_and_consumer_cancels(): + """Stops when the stream queue is back-pressured and the consumer cancels + + The JavaScript version cancels by calling return() on the never pulled + iterator; the asyncio analog of such a concurrent consumer cancellation + is cancelling the task that awaits a pending anext(), so the stream is + additionally held open with a pending first item to keep the pull + pending while the producer is parked on the back-pressured queue. + """ + document = parse("{ scalarList @stream(initialCount: 0) }") + + item_future: Future[str] = Future() + reached_capacity = Event() + + class Source: + def __init__(self): + self.count = 0 + self.close_calls = 0 + + def __iter__(self): + return self + + def __next__(self): + self.count += 1 + if self.count == 1: + return item_future # pending head item + if self.count == 101: + reached_capacity.set() + if self.count > 101: + raise StopIteration + return str(self.count) + + def close(self): + self.close_calls += 1 # pragma: no cover + + source = Source() + + result = experimental_execute_incrementally( + cancellation_schema, + document, + {"scalarList": lambda _info: source}, + enable_early_execution=True, + ) + assert isinstance(result, ExperimentalIncrementalExecutionResults) + iterator = result.subsequent_results + + await reached_capacity.wait() + for _ in range(2): # let the producer park on the back-pressured queue + await sleep(0) + assert source.count == 101 + + next_task = ensure_future(anext(iterator)) + for _ in range(2): + await sleep(0) + # the pull stays pending on the first item while the producer is parked + assert not next_task.done() + + next_task.cancel() + with pytest.raises(CancelledError): + await next_task + + # the pending head item was cancelled with the parked producer + assert item_future.cancelled() + # the cancelled sync iterator was drained, but not closed + assert source.count == 102 + assert source.close_calls == 0 + async def stops_streaming_when_pending_stream_item_is_cancelled(): """Stops streaming when a pending stream item resolves after cancellation diff --git a/tests/execution/incremental/test_stream_item_queue.py b/tests/execution/incremental/test_stream_item_queue.py new file mode 100644 index 00000000..308df393 --- /dev/null +++ b/tests/execution/incremental/test_stream_item_queue.py @@ -0,0 +1,360 @@ +"""Unit tests for the stream item queue. + +These tests cover the asyncio-specific seams of the stream item queue that +are not (or not deterministically) reachable via the defer/stream test +suites, especially the batching of pending entries and the abort paths. +""" + +from __future__ import annotations + +from asyncio import ( + CancelledError, + Event, + Future, + ensure_future, + get_running_loop, + sleep, +) +from typing import TYPE_CHECKING, Any + +import pytest + +from graphql.execution.incremental import StreamItemQueue, WorkResult +from graphql.pyutils import is_awaitable + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Sequence + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +def make_queue( + produce: Callable[[StreamItemQueue], Any], + on_abort: Callable[[BaseException | None], Awaitable[None] | None] | None = None, + eager: bool = False, + capacity: int = 100, +) -> StreamItemQueue: + return StreamItemQueue(produce, on_abort, eager=eager, capacity=capacity) + + +async def collect_batches(queue: StreamItemQueue) -> list[Sequence[WorkResult]]: + return [batch async for batch in queue.batches()] + + +def describe_stream_item_queue(): + async def delivers_settled_results_in_one_batch(): + async def produce(queue: StreamItemQueue) -> None: + await queue.push(WorkResult(1)) + await queue.push(WorkResult(2)) + + queue = make_queue(produce) + assert not queue.is_stopped() + batches = await collect_batches(queue) + assert batches == [[WorkResult(1), WorkResult(2)]] + assert queue.is_stopped() + + async def delivers_batches_in_order_with_pending_futures(): + loop_futures: list[Future[WorkResult]] = [] + + async def produce(queue: StreamItemQueue) -> None: + loop = get_running_loop() + first: Future[WorkResult] = loop.create_future() + second: Future[WorkResult] = loop.create_future() + loop_futures.extend([first, second]) + await queue.push(first) + await queue.push(second) + await queue.push(WorkResult(3)) + # settle out of order: the batches must stay in source order + second.set_result(WorkResult(2)) + first.set_result(WorkResult(1)) + + queue = make_queue(produce) + batches = await collect_batches(queue) + assert batches in ( + [[WorkResult(1), WorkResult(2), WorkResult(3)]], + [[WorkResult(1)], [WorkResult(2), WorkResult(3)]], + ) + + async def holds_back_pending_futures_at_the_end_of_a_batch(): + release = Event() + + async def produce(queue: StreamItemQueue) -> None: + await queue.push(WorkResult(1)) + loop = get_running_loop() + future: Future[WorkResult] = loop.create_future() + await queue.push(future) + + async def settle_later() -> None: + await release.wait() + future.set_result(WorkResult(2)) + + task = ensure_future(settle_later()) + release.set() + await task + + queue = make_queue(produce) + batches = await collect_batches(queue) + assert batches == [[WorkResult(1)], [WorkResult(2)]] + + async def raises_stream_failure_after_delivering_settled_results(): + error = RuntimeError("stream failed") + + async def produce(queue: StreamItemQueue) -> None: + await queue.push(WorkResult(1)) + raise error + + queue = make_queue(produce) + batches: list[Sequence[WorkResult]] = [] + + async def consume() -> None: + async for batch in queue.batches(): + batches.append(batch) # noqa: PERF401 + + with pytest.raises(RuntimeError, match="stream failed"): + await consume() + assert batches == [[WorkResult(1)]] + assert not queue.is_stopped() + + async def runs_cleanup_before_delivering_a_stream_failure(): + cleaned_up: list[BaseException | None] = [] + + async def produce(_queue: StreamItemQueue) -> None: + raise RuntimeError("stream failed") + + async def async_cleanup() -> None: + cleaned_up.append(None) + + def on_abort(reason: BaseException | None) -> Awaitable[None]: + cleaned_up.append(reason) + return async_cleanup() + + queue = make_queue(produce, on_abort) + with pytest.raises(RuntimeError, match="stream failed"): + await collect_batches(queue) + assert len(cleaned_up) == 2 + assert isinstance(cleaned_up[0], RuntimeError) + + async def delivers_rejection_of_a_pending_head_after_cleanup(): + aborted: list[BaseException | None] = [] + + async def produce(queue: StreamItemQueue) -> None: + loop = get_running_loop() + failing: Future[WorkResult] = loop.create_future() + failing.set_exception(RuntimeError("item failed")) + await queue.push(failing) + await Event().wait() # never finishes normally + + queue = make_queue(produce, aborted.append) + with pytest.raises(RuntimeError, match="item failed"): + await collect_batches(queue) + assert len(aborted) == 1 + assert queue.abort() is None # already aborted + + async def holds_back_a_rejected_future_at_the_end_of_a_batch(): + async def produce(queue: StreamItemQueue) -> None: + await queue.push(WorkResult(1)) + loop = get_running_loop() + failing: Future[WorkResult] = loop.create_future() + failing.set_exception(RuntimeError("item failed")) + await queue.push(failing) + + queue = make_queue(produce) + batches: list[Sequence[WorkResult]] = [] + + async def consume() -> None: + async for batch in queue.batches(): + batches.append(batch) # noqa: PERF401 + + with pytest.raises(RuntimeError, match="item failed"): + await consume() + assert batches == [[WorkResult(1)]] + + async def starts_the_producer_eagerly_when_a_loop_is_running(): + started = Event() + + async def produce(queue: StreamItemQueue) -> None: + started.set() + await queue.push(WorkResult(1)) + + queue = make_queue(produce, eager=True) + await sleep(0) + assert started.is_set() # started before consumption + batches = await collect_batches(queue) + assert batches == [[WorkResult(1)]] + + def starts_the_producer_lazily_without_a_running_loop(): + async def produce(_queue: StreamItemQueue) -> None: + return # pragma: no cover + + queue = make_queue(produce, eager=True) # no loop running here + assert queue._producer_task is None # noqa: SLF001 + + async def aborting_an_unconsumed_queue_runs_the_cleanup_synchronously(): + aborted: list[BaseException | None] = [] + + async def produce(_queue: StreamItemQueue) -> None: + return # pragma: no cover + + queue = make_queue(produce, aborted.append) + reason = RuntimeError("abort") + assert queue.abort(reason) is None + assert aborted == [reason] + + async def aborting_a_started_queue_cancels_the_producer(): + aborted: list[BaseException | None] = [] + parked = Event() + + async def produce(_queue: StreamItemQueue) -> None: + parked.set() + await Event().wait() # park forever + + queue = make_queue(produce, aborted.append, eager=True) + await parked.wait() + cleanup = queue.abort() + assert is_awaitable(cleanup) + await cleanup + assert aborted == [None] + producer_task = queue._producer_task # noqa: SLF001 + assert producer_task is not None + assert producer_task.cancelled() + + async def aborting_a_started_queue_cancels_pending_item_futures(): + pending: list[Future[WorkResult]] = [] + + async def parked_item() -> WorkResult: + await Event().wait() # park forever + return WorkResult(None) # pragma: no cover + + async def produce(queue: StreamItemQueue) -> None: + future = ensure_future(parked_item()) + pending.append(future) + await queue.push(future) + + queue = make_queue(produce, eager=True) + while not pending: # noqa: ASYNC110 + await sleep(0) + await sleep(0) # let the producer finish normally + cleanup = queue.abort() + assert is_awaitable(cleanup) + await cleanup + assert pending[0].cancelled() + + async def aborting_an_unconsumed_queue_without_cleanup_callback(): + async def produce(_queue: StreamItemQueue) -> None: + return # pragma: no cover + + queue = make_queue(produce) + assert queue.abort() is None + + async def awaiting_cleanup_settles_pending_item_futures(): + pushed = Event() + + async def parked_item() -> WorkResult: + await Event().wait() # park forever + return WorkResult(None) # pragma: no cover + + async def produce(queue: StreamItemQueue) -> None: + await queue.push(ensure_future(parked_item())) + pushed.set() + await Event().wait() # park forever + + queue = make_queue(produce, eager=True) + await pushed.wait() + cleanup = queue.abort() + assert is_awaitable(cleanup) + await cleanup + + async def aborting_after_normal_finish_skips_the_cleanup_callback(): + aborted: list[BaseException | None] = [] + + async def produce(queue: StreamItemQueue) -> None: + await queue.push(WorkResult(1)) + + queue = make_queue(produce, aborted.append) + batches = await collect_batches(queue) + assert batches == [[WorkResult(1)]] + assert queue.abort() is None + assert aborted == [] # the source finished and must not be cleaned up + + async def awaiting_cleanup_awaits_the_abort_callback(): + cleanup_done = Event() + parked = Event() + + async def produce(_queue: StreamItemQueue) -> None: + parked.set() + await Event().wait() # park forever + + async def async_cleanup() -> None: + await sleep(0) + cleanup_done.set() + + def on_abort(_reason: BaseException | None) -> Awaitable[None]: + return async_cleanup() + + queue = make_queue(produce, on_abort, eager=True) + await parked.wait() + cleanup = queue.abort() + assert is_awaitable(cleanup) + await cleanup + assert cleanup_done.is_set() + + async def consumer_cancellation_does_not_settle_pending_items(): + pushed = Event() + + async def produce(queue: StreamItemQueue) -> None: + future: Future[WorkResult] = get_running_loop().create_future() + await queue.push(future) + pushed.set() + await Event().wait() # park forever + + queue = make_queue(produce) + + async def consume() -> None: + async for _batch in queue.batches(): # pragma: no cover + pass + + consumer = ensure_future(consume()) + await pushed.wait() + consumer.cancel() + with pytest.raises(CancelledError): + await consumer + cleanup = queue.abort() + assert is_awaitable(cleanup) + await cleanup + + async def aborting_releases_a_producer_parked_on_the_end_marker(): + aborted: list[BaseException | None] = [] + + async def produce(queue: StreamItemQueue) -> None: + await queue.push(WorkResult(1)) + + queue = make_queue(produce, aborted.append, eager=True, capacity=1) + for _ in range(3): # let the producer park on the end marker + await sleep(0) + cleanup = queue.abort() + assert is_awaitable(cleanup) + await cleanup + assert aborted == [] # the source finished and must not be cleaned up + assert queue.abort() is None # the producer has already been released + + async def aborting_again_releases_a_producer_parked_on_a_failure_entry(): + aborted: list[BaseException | None] = [] + + async def produce(queue: StreamItemQueue) -> None: + await queue.push(WorkResult(1)) + raise RuntimeError("boom") + + queue = make_queue(produce, aborted.append, eager=True, capacity=1) + for _ in range(3): # let the producer park on the failure entry + await sleep(0) + assert len(aborted) == 1 # the producer has already run the cleanup + cleanup = queue.abort() + assert is_awaitable(cleanup) + await cleanup + assert len(aborted) == 1 # the cleanup must not run again + assert queue.abort() is None # the producer has already been released diff --git a/tests/execution/incremental/test_work_queue.py b/tests/execution/incremental/test_work_queue.py new file mode 100644 index 00000000..fc9e0010 --- /dev/null +++ b/tests/execution/incremental/test_work_queue.py @@ -0,0 +1,1243 @@ +from __future__ import annotations + +from asyncio import CancelledError, Event, Future, ensure_future, sleep +from typing import TYPE_CHECKING, Any + +import pytest + +from graphql.execution.incremental import Computation +from graphql.execution.incremental.work_queue import ( + GroupFailureEvent, + GroupSuccessEvent, + GroupValuesEvent, + StreamFailureEvent, + StreamQueue, + StreamSuccessEvent, + StreamValuesEvent, + Work, + WorkQueue, + WorkQueueTerminationEvent, + WorkResult, + WorkTask, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Awaitable, Callable, Sequence + +pytestmark = pytest.mark.anyio + + +class Spy: + """Wrap a function, counting its calls.""" + + def __init__(self, fn: Callable[..., Any]) -> None: + self.fn = fn + self.call_count = 0 + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self.call_count += 1 + return self.fn(*args, **kwargs) + + +class Group: + """A test group with an optional parent and a name for debugging.""" + + def __init__(self, parent: Group | None = None, name: str = "group") -> None: + self.parent = parent + self.name = name + + def __repr__(self) -> str: + return f"" + + +class FakeStreamQueue: + """A fake stream queue yielding predefined batches of stream items.""" + + def __init__( + self, + item_batches: Sequence[Sequence[WorkResult]] = (), + error: BaseException | None = None, + on_abort: Callable[[BaseException | None], Awaitable[None] | None] + | None = None, + never_stops: bool = False, + stops_early: bool = False, + ) -> None: + self.item_batches = item_batches + self.error = error + self.on_abort = on_abort + self.never_stops = never_stops + self.stops_early = stops_early + self.stopped = False + self.aborted_with: list[BaseException | None] = [] + + async def batches(self) -> AsyncIterator[Sequence[WorkResult]]: + item_batches = self.item_batches + last_index = len(item_batches) - 1 + for index, item_batch in enumerate(item_batches): + if index == last_index and self.stops_early: + # known to have stopped while delivering the last batch, + # like the real stream item queue for fast sources + self.stopped = True + yield item_batch + if self.error: + raise self.error + if self.never_stops: + # cancelled while parked, possibly before this step runs + await Event().wait() # pragma: no cover + self.stopped = True + + def is_stopped(self) -> bool: + return self.stopped + + def abort(self, reason: BaseException | None = None) -> Awaitable[None] | None: + self.aborted_with.append(reason) + if self.on_abort: + return self.on_abort(reason) + return None + + +class Stream: + """A test stream with a name for debugging.""" + + def __init__(self, queue: FakeStreamQueue, name: str = "stream") -> None: + self.queue: StreamQueue = queue + self.name = name + + def __repr__(self) -> str: + return f"" + + +def make_task( + groups: Sequence[Group], + value_or_fn: Any, + work: Work | None = None, +) -> WorkTask: + if callable(value_or_fn): + return WorkTask(groups, Computation(value_or_fn)) + return WorkTask(groups, Computation(lambda: WorkResult(value_or_fn, work))) + + +def stream_from( + items: Sequence[WorkResult], + error: BaseException | None = None, + batched: bool = False, + name: str = "stream", +) -> Stream: + item_batches = [items] if batched else [[item] for item in items] + return Stream(FakeStreamQueue(item_batches, error=error), name) + + +async def collect_work_run(work: Work) -> tuple[Any, Any, list[Any]]: + events: list[Any] = [] + work_queue = WorkQueue(work) + async for batch in work_queue.events(): + events.extend(batch) + return work_queue.initial_groups, work_queue.initial_streams, events + + +def describe_work_queue(): + async def runs_parent_and_child_groups_sequentially(): + root = Group(name="root") + child = Group(root, name="child") + + child_ran_spy = Spy(lambda: WorkResult("child")) + child_task = make_task([child], child_ran_spy) + root_task = make_task([root], "root", Work(groups=[child], tasks=[child_task])) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[root], tasks=[root_task]) + ) + + assert initial_groups == [root] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(root, ["root"]), + GroupSuccessEvent(root, [child], []), + GroupValuesEvent(child, ["child"]), + GroupSuccessEvent(child, [], []), + WorkQueueTerminationEvent(), + ] + assert child_ran_spy.call_count == 1 + + async def can_handle_child_groups_passed_prior_to_parents(): + root = Group(name="root") + child = Group(root, name="child") + + child_task = make_task([child], "child", Work()) + root_task = make_task([root], "root", Work()) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[child, root], tasks=[root_task, child_task]) + ) + + assert initial_groups == [root] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(root, ["root"]), + GroupSuccessEvent(root, [child], []), + GroupValuesEvent(child, ["child"]), + GroupSuccessEvent(child, [], []), + WorkQueueTerminationEvent(), + ] + + async def propagates_task_failures_and_skips_descendant_groups(): + root = Group(name="root") + child = Group(root, name="child") + grandchild = Group(child, name="grandchild") + grandchild_ran_spy = Spy(lambda: WorkResult("grandchild")) + grandchild_task = make_task([grandchild], grandchild_ran_spy) + + boom = RuntimeError("boom") + + def fail() -> WorkResult: + raise boom + + child_ran_spy = Spy(fail) + failing_child_task = make_task([child], child_ran_spy) + + root_task = make_task( + [root], + "root", + Work( + groups=[child, grandchild], tasks=[failing_child_task, grandchild_task] + ), + ) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[root], tasks=[root_task]) + ) + + assert initial_groups == [root] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(root, ["root"]), + GroupSuccessEvent(root, [child], []), + GroupFailureEvent(child, boom), + WorkQueueTerminationEvent(), + ] + assert grandchild_ran_spy.call_count == 0 + assert child_ran_spy.call_count == 1 + + async def integrates_work_object_returned_by_task(): + root = Group(name="root") + child = Group(root, name="child") + + child_task = make_task([child], lambda: WorkResult("child")) + root_task = make_task( + [root], + lambda: WorkResult("root", Work(groups=[child], tasks=[child_task])), + ) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[root], tasks=[root_task]) + ) + + assert initial_groups == [root] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(root, ["root"]), + GroupSuccessEvent(root, [child], []), + GroupValuesEvent(child, ["child"]), + GroupSuccessEvent(child, [], []), + WorkQueueTerminationEvent(), + ] + + async def purges_shared_tasks_so_sibling_groups_finish_without_rerunning_work(): + group_a = Group(name="a") + group_b = Group(name="b") + + shared_task = make_task([group_a, group_b], "shared") + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group_a, group_b], tasks=[shared_task]) + ) + + assert initial_groups == [group_a, group_b] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(group_a, ["shared"]), + GroupSuccessEvent(group_a, [], []), + GroupSuccessEvent(group_b, [], []), + WorkQueueTerminationEvent(), + ] + + async def ignores_task_emitted_groups_without_a_valid_parent(): + root = Group(name="root") + orphan_group = Group(name="orphan") + + task = make_task([root], "root", Work(groups=[orphan_group])) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[root], tasks=[task]) + ) + + assert initial_groups == [root] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(root, ["root"]), + GroupSuccessEvent(root, [], []), + WorkQueueTerminationEvent(), + ] + + async def skips_child_groups_with_only_completed_tasks_when_parent_finishes_later(): + parent1 = Group(name="parent1") + parent2 = Group(name="parent2") + child1 = Group(parent1, name="child1") + child2 = Group(parent2, name="child2") + + parent1_task = make_task([parent1], "parent1") + + async def slow_parent2() -> WorkResult: + await sleep(0) + return WorkResult("parent2-slow") + + slow_parent2_task = make_task([parent2], slow_parent2) + shared_child_task = make_task([child1, child2], "child-shared") + + async def slow_child1_follow_up() -> WorkResult: + await sleep(0) + await sleep(0) + await sleep(0) + return WorkResult("child1-slow") + + slow_child1_follow_up_task = make_task([child1], slow_child1_follow_up) + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[parent1, parent2, child1, child2], + tasks=[ + parent1_task, + slow_parent2_task, + shared_child_task, + slow_child1_follow_up_task, + ], + ) + ) + + assert initial_groups == [parent1, parent2] + assert initial_streams == [] + # Note: the upstream test expects these events for `child2`, but its + # deep equality cannot distinguish the structurally identical groups; + # the events are actually emitted for `child1` (`child2` is pruned). + assert events == [ + GroupValuesEvent(parent1, ["parent1"]), + GroupSuccessEvent(parent1, [child1], []), + GroupValuesEvent(parent2, ["parent2-slow"]), + GroupSuccessEvent(parent2, [], []), + GroupValuesEvent(child1, ["child-shared", "child1-slow"]), + GroupSuccessEvent(child1, [], []), + WorkQueueTerminationEvent(), + ] + + async def skips_promoted_child_groups_that_already_completed_shared_tasks(): + parent = Group(name="parent") + child = Group(parent, name="child") + + async def slow_parent() -> WorkResult: + await sleep(0) + await sleep(0) + return WorkResult("parent") + + parent_task = make_task([parent], slow_parent) + shared_task = make_task([parent, child], "shared") + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[parent, child], tasks=[parent_task, shared_task]) + ) + + assert initial_groups == [parent] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(parent, ["parent", "shared"]), + GroupSuccessEvent(parent, [], []), + WorkQueueTerminationEvent(), + ] + + async def skips_child_groups_with_shared_tasks_completed_by_a_parent(): + parent = Group(name="parent") + child = Group(parent, name="child") + other_root = Group(name="other_root") + + async def slow_parent() -> WorkResult: + await sleep(0) + await sleep(0) + return WorkResult("parent") + + parent_task = make_task([parent], slow_parent) + shared_task = make_task([child, other_root], "shared") + + async def slow_other_root() -> WorkResult: + await sleep(0) + await sleep(0) + await sleep(0) + await sleep(0) + return WorkResult("other-root") + + slow_other_root_task = make_task([other_root], slow_other_root) + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[parent, other_root, child], + tasks=[parent_task, shared_task, slow_other_root_task], + ) + ) + + assert initial_groups == [parent, other_root] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(parent, ["parent"]), + GroupSuccessEvent(parent, [], []), + GroupValuesEvent(other_root, ["shared", "other-root"]), + GroupSuccessEvent(other_root, [], []), + WorkQueueTerminationEvent(), + ] + + async def does_not_promote_child_groups_that_only_share_work_with_the_parent(): + parent = Group(name="parent") + child = Group(parent, name="child") + + shared_task = make_task([parent, child], "shared") + + async def slow_parent_only() -> WorkResult: + await sleep(0) + return WorkResult("parent-only") + + parent_only_task = make_task([parent], slow_parent_only) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[parent, child], tasks=[shared_task, parent_only_task]) + ) + + assert initial_groups == [parent] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(parent, ["shared", "parent-only"]), + GroupSuccessEvent(parent, [], []), + WorkQueueTerminationEvent(), + ] + + async def ignores_work_returned_by_tasks_whose_groups_already_failed(): + group = Group(name="group") + late_group = Group(name="late_group") + + fail_early = RuntimeError("fail early") + + def fail() -> WorkResult: + raise fail_early + + failing = make_task([group], fail) + + async def slow() -> WorkResult: + await sleep(0) + return WorkResult("late", Work(groups=[late_group])) + + slow_task = make_task([group], slow) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[failing, slow_task]) + ) + + assert initial_groups == [group] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(group, fail_early), + WorkQueueTerminationEvent(), + ] + + async def defers_task_emitted_streams_until_the_parent_group_succeeds(): + parent = Group(name="parent") + deferred_stream = stream_from([WorkResult(7)], name="deferred") + + parent_task = make_task([parent], "parent", Work(streams=[deferred_stream])) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[parent], tasks=[parent_task]) + ) + + assert initial_groups == [parent] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(parent, ["parent"]), + GroupSuccessEvent(parent, [], [deferred_stream]), + StreamValuesEvent(deferred_stream, [7], [], []), + StreamSuccessEvent(deferred_stream), + WorkQueueTerminationEvent(), + ] + + async def only_promotes_shared_streams_after_all_parent_groups_finish(): + group_a = Group(name="a") + group_b = Group(name="b") + shared_stream = stream_from([WorkResult(1)], name="shared") + + shared_task = make_task( + [group_a, group_b], "shared", Work(streams=[shared_stream]) + ) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group_a, group_b], tasks=[shared_task]) + ) + + assert initial_groups == [group_a, group_b] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(group_a, ["shared"]), + GroupSuccessEvent(group_a, [], [shared_stream]), + GroupSuccessEvent(group_b, [], []), + StreamValuesEvent(shared_stream, [1], [], []), + StreamSuccessEvent(shared_stream), + WorkQueueTerminationEvent(), + ] + + async def starts_a_shared_stream_once_even_when_the_second_parent_is_slower(): + group_a = Group(name="a") + group_b = Group(name="b") + shared_stream = stream_from([WorkResult(5)], name="shared") + + shared_task = make_task( + [group_a, group_b], "shared", Work(streams=[shared_stream]) + ) + fast_task_a = make_task([group_a], "A-only") + + async def slow_b() -> WorkResult: + await sleep(0) + await sleep(0) + await sleep(0) + return WorkResult("B-slow") + + slow_task_b = make_task([group_b], slow_b) + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[group_a, group_b], tasks=[shared_task, fast_task_a, slow_task_b] + ) + ) + + assert initial_groups == [group_a, group_b] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(group_a, ["shared", "A-only"]), + GroupSuccessEvent(group_a, [], [shared_stream]), + StreamValuesEvent(shared_stream, [5], [], []), + StreamSuccessEvent(shared_stream), + GroupValuesEvent(group_b, ["B-slow"]), + GroupSuccessEvent(group_b, [], []), + WorkQueueTerminationEvent(), + ] + + async def does_not_promote_a_child_stream_if_the_parent_fails(): + group = Group(name="group") + stream = stream_from([WorkResult(99)]) + + task = make_task([group], "task", Work(streams=[stream])) + boom = RuntimeError("boom") + + def fail() -> WorkResult: + raise boom + + failing_task = make_task([group], fail) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[task, failing_task]) + ) + + assert initial_groups == [group] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(group, boom), + WorkQueueTerminationEvent(), + ] + + async def promotes_a_stream_with_multiple_parents_when_only_a_single_parent_fails(): + group_a = Group(name="a") + group_b = Group(name="b") + shared_stream = stream_from([WorkResult(99)], name="shared") + + shared_task = make_task( + [group_a, group_b], "shared", Work(streams=[shared_stream]) + ) + boom = RuntimeError("boom") + + def fail() -> WorkResult: + raise boom + + failing_task = make_task([group_a], fail) + + async def slow_b() -> WorkResult: + await sleep(0) + return WorkResult("B-resolved") + + slow_task_b = make_task([group_b], slow_b) + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[group_a, group_b], + tasks=[shared_task, failing_task, slow_task_b], + ) + ) + + assert initial_groups == [group_a, group_b] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(group_a, boom), + GroupValuesEvent(group_b, ["shared", "B-resolved"]), + GroupSuccessEvent(group_b, [], [shared_stream]), + StreamValuesEvent(shared_stream, [99], [], []), + StreamSuccessEvent(shared_stream), + WorkQueueTerminationEvent(), + ] + + async def emits_stream_items_followed_by_success(): + stream = stream_from([WorkResult(1), WorkResult(2), WorkResult(3)]) + + initial_groups, initial_streams, events = await collect_work_run( + Work(streams=[stream]) + ) + + assert initial_groups == [] + assert initial_streams == [stream] + assert events == [ + StreamValuesEvent(stream, [1], [], []), + StreamValuesEvent(stream, [2], [], []), + StreamValuesEvent(stream, [3], [], []), + StreamSuccessEvent(stream), + WorkQueueTerminationEvent(), + ] + + async def handles_batched_stream_items(): + spawned = Group(name="spawned") + spawned_task = make_task([spawned], "spawned-from-stream") + + stream = stream_from( + [ + WorkResult(1), + WorkResult(2, Work(groups=[spawned], tasks=[spawned_task])), + ], + batched=True, + ) + + initial_groups, initial_streams, events = await collect_work_run( + Work(streams=[stream]) + ) + + assert initial_groups == [] + assert initial_streams == [stream] + assert events == [ + StreamValuesEvent(stream, [1, 2], [spawned], []), + GroupValuesEvent(spawned, ["spawned-from-stream"]), + GroupSuccessEvent(spawned, [], []), + StreamSuccessEvent(stream), + WorkQueueTerminationEvent(), + ] + + async def emits_stream_failure_when_the_iterator_throws(): + broken_stream_error = RuntimeError("broken stream") + failing_stream = stream_from([WorkResult(42)], error=broken_stream_error) + + initial_groups, initial_streams, events = await collect_work_run( + Work(streams=[failing_stream]) + ) + + assert initial_groups == [] + assert initial_streams == [failing_stream] + assert events == [ + StreamValuesEvent(failing_stream, [42], [], []), + StreamFailureEvent(failing_stream, broken_stream_error), + WorkQueueTerminationEvent(), + ] + + async def emits_stream_success_in_a_later_payload_when_stream_is_slow_to_stop(): + proceed: Future[None] = Future() + + class SlowToStopQueue(FakeStreamQueue): + async def batches(self) -> AsyncIterator[Sequence[WorkResult]]: + yield [WorkResult(1)] + await proceed + + stream = Stream(SlowToStopQueue()) + + events: list[Any] = [] + work_queue = WorkQueue(Work(streams=[stream])) + iterator = work_queue.events() + events.extend(await anext(iterator)) + proceed.set_result(None) + async for batch in iterator: + events.extend(batch) + + assert events == [ + StreamValuesEvent(stream, [1], [], []), + StreamSuccessEvent(stream), + WorkQueueTerminationEvent(), + ] + + async def emits_late_root_groups_and_streams_triggered_from_stream(): + initial = Group(name="initial") + late_group = Group(name="late_group") + + late_task = make_task([late_group], "late") + secondary_stream = stream_from([WorkResult(9)], name="secondary") + trigger_stream = stream_from( + [ + WorkResult( + 0, + Work( + groups=[late_group], + tasks=[late_task], + streams=[secondary_stream], + ), + ) + ], + name="trigger", + ) + + initial_task = make_task([initial], "initial") + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[initial], tasks=[initial_task], streams=[trigger_stream]) + ) + + assert initial_groups == [initial] + assert initial_streams == [trigger_stream] + assert events == [ + GroupValuesEvent(initial, ["initial"]), + GroupSuccessEvent(initial, [], []), + StreamValuesEvent(trigger_stream, [0], [late_group], [secondary_stream]), + GroupValuesEvent(late_group, ["late"]), + GroupSuccessEvent(late_group, [], []), + StreamValuesEvent(secondary_stream, [9], [], []), + StreamSuccessEvent(trigger_stream), + StreamSuccessEvent(secondary_stream), + WorkQueueTerminationEvent(), + ] + + async def handles_tasks_that_are_started_manually_before_they_complete(): + group = Group(name="group") + + async def primed() -> WorkResult: + await sleep(0) + return WorkResult("primed") + + computation: Computation[WorkResult] = Computation(primed) + primed_task = WorkTask([group], computation) + + computation.prime() + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[primed_task]) + ) + + assert initial_groups == [group] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(group, ["primed"]), + GroupSuccessEvent(group, [], []), + WorkQueueTerminationEvent(), + ] + + async def propagates_failures_for_tasks_started_manually(): + group = Group(name="group") + primed_failure = RuntimeError("primed failure") + + async def primed() -> WorkResult: + await sleep(0) + raise primed_failure + + computation: Computation[WorkResult] = Computation(primed) + primed_task = WorkTask([group], computation) + + computation.prime() + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[primed_task]) + ) + + assert initial_groups == [group] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(group, primed_failure), + WorkQueueTerminationEvent(), + ] + + async def skips_groups_with_no_tasks_and_promotes_descendants(): + root = Group(name="root") + empty_parent = Group(root, name="empty_parent") + leaf = Group(empty_parent, name="leaf") + + leaf_task = make_task([leaf], "leaf") + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[root, empty_parent, leaf], tasks=[leaf_task]) + ) + + assert initial_groups == [leaf] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(leaf, ["leaf"]), + GroupSuccessEvent(leaf, [], []), + WorkQueueTerminationEvent(), + ] + + async def handles_tasks_that_are_already_settled_before_being_queued(): + group = Group(name="group") + computation: Computation[WorkResult] = Computation(lambda: WorkResult("eager")) + eager_task = WorkTask([group], computation) + + computation.prime() + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[eager_task]) + ) + + assert initial_groups == [group] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(group, ["eager"]), + GroupSuccessEvent(group, [], []), + WorkQueueTerminationEvent(), + ] + + async def cancels_nested_work_when_the_work_queue_is_cancelled(): + root = Group(name="root") + child = Group(root, name="child") + root_stream = stream_from([WorkResult(1)], name="root_stream") + + child_stream_cancelled = False + child_stream_cleanup: Future[None] = Future() + + def cancel_child_stream(_reason: BaseException | None) -> Future[None]: + nonlocal child_stream_cancelled + child_stream_cancelled = True + return child_stream_cleanup + + child_stream = Stream( + FakeStreamQueue(never_stops=True, on_abort=cancel_child_stream), + name="child_stream", + ) + + child_task_cancelled = False + child_task_cleanup: Future[None] = Future() + + def cancel_child_task(_reason: BaseException | None) -> Future[None]: + nonlocal child_task_cancelled + child_task_cancelled = True + return child_task_cleanup + + async def never() -> WorkResult: + # cancelled while parked, possibly before this step runs + await Event().wait() # pragma: no cover + raise AssertionError("never resolves") # pragma: no cover + + child_task = WorkTask([root], Computation(never, cancel_child_task)) + + root_task = make_task( + [root], lambda: WorkResult("root", Work(streams=[child_stream])) + ) + + work_queue = WorkQueue( + Work( + groups=[root, child], + tasks=[root_task, child_task], + streams=[root_stream], + ) + ) + + assert work_queue.initial_groups == [root] + assert work_queue.initial_streams == [root_stream] + iterator = work_queue.events() + assert await anext(iterator) == [ + StreamValuesEvent(root_stream, [1], [], []), + ] + + cancelled = ensure_future(work_queue.cancel()) + await sleep(0) + assert child_task_cancelled is True + assert child_stream_cancelled is True + assert not cancelled.done() # cancellation awaits async cleanup + + child_stream_cleanup.set_result(None) + child_task_cleanup.set_result(None) + await cancelled + + with pytest.raises(StopAsyncIteration): + await anext(iterator) + + async def forwards_cancellation_reason_when_cancelling_nested_work(): + root = Group(name="root") + root_stream = stream_from([WorkResult(1)], name="root_stream") + abort_reason = RuntimeError("Abort nested work") + + child_stream_cancel_reasons: list[BaseException | None] = [] + child_stream = Stream( + FakeStreamQueue( + never_stops=True, on_abort=child_stream_cancel_reasons.append + ), + name="child_stream", + ) + + child_task_cancel_reasons: list[BaseException | None] = [] + + async def never() -> WorkResult: + # cancelled while parked, possibly before this step runs + await Event().wait() # pragma: no cover + raise AssertionError("never resolves") # pragma: no cover + + child_task = WorkTask( + [root], Computation(never, child_task_cancel_reasons.append) + ) + + root_task = make_task( + [root], lambda: WorkResult("root", Work(streams=[child_stream])) + ) + + work_queue = WorkQueue( + Work(groups=[root], tasks=[root_task, child_task], streams=[root_stream]) + ) + + iterator = work_queue.events() + await anext(iterator) + + await work_queue.cancel(abort_reason) + + assert child_task_cancel_reasons == [abort_reason] + assert child_stream_cancel_reasons == [abort_reason] + + async def ends_iteration_when_cancelled_while_waiting_for_events(): + group = Group(name="group") + + async def never() -> WorkResult: + # cancelled while parked, possibly before this step runs + await Event().wait() # pragma: no cover + raise AssertionError("never resolves") # pragma: no cover + + never_task = make_task([group], never) + work_queue = WorkQueue(Work(groups=[group], tasks=[never_task])) + + events: list[Any] = [] + + async def consume() -> None: + async for batch in work_queue.events(): + events.extend(batch) # pragma: no cover + + consumer = ensure_future(consume()) + await sleep(0) # park the consumer on the event channel + await work_queue.cancel() + await consumer # the stop sentinel ends the iteration + + assert events == [] + + async def aborts_unstarted_tasks_of_pending_child_groups_when_cancelled(): + parent = Group(name="parent") + child = Group(parent, name="child") + + async def never() -> WorkResult: + # cancelled while parked, possibly before this step runs + await Event().wait() # pragma: no cover + raise AssertionError("never resolves") # pragma: no cover + + never_task = make_task([parent], never) + run_spy = Spy(lambda: WorkResult("child")) + lazy_child_task = make_task([child], run_spy) + + work_queue = WorkQueue( + Work(groups=[parent, child], tasks=[never_task, lazy_child_task]) + ) + iterator = work_queue.events() + consumer = ensure_future(anext(iterator)) + await sleep(0) + await work_queue.cancel() + with pytest.raises(StopAsyncIteration): + await consumer + + assert run_spy.call_count == 0 # poisoned before running + with pytest.raises(CancelledError): + lazy_child_task.computation.result() + + async def ignores_tasks_and_streams_emitted_for_already_failed_groups(): + group = Group(name="group") + + fail_early = RuntimeError("fail early") + + def fail() -> WorkResult: + raise fail_early + + failing = make_task([group], fail) + + run_spy = Spy(lambda: WorkResult("more")) + late_stream = stream_from([WorkResult(1)], name="late") + + async def slow() -> WorkResult: + await sleep(0) + return WorkResult( + "late", + Work(tasks=[make_task([group], run_spy)], streams=[late_stream]), + ) + + slow_task = make_task([group], slow) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[failing, slow_task]) + ) + + assert initial_groups == [group] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(group, fail_early), + WorkQueueTerminationEvent(), + ] + assert run_spy.call_count == 0 + + async def promotes_multiple_streams_from_the_same_task_together(): + parent = Group(name="parent") + stream1 = stream_from([WorkResult(1)], name="stream1") + stream2 = stream_from([WorkResult(2)], name="stream2") + + parent_task = make_task([parent], "parent", Work(streams=[stream1, stream2])) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[parent], tasks=[parent_task]) + ) + + assert initial_groups == [parent] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(parent, ["parent"]), + GroupSuccessEvent(parent, [], [stream1, stream2]), + StreamValuesEvent(stream1, [1], [], []), + StreamValuesEvent(stream2, [2], [], []), + StreamSuccessEvent(stream1), + StreamSuccessEvent(stream2), + WorkQueueTerminationEvent(), + ] + + async def finishes_shared_tasks_that_settle_after_one_of_their_groups_failed(): + group_a = Group(name="a") + group_b = Group(name="b") + + boom = RuntimeError("boom") + + def fail() -> WorkResult: + raise boom + + failing_task = make_task([group_a], fail) + + async def slow_shared() -> WorkResult: + await sleep(0) + await sleep(0) + return WorkResult("shared-slow") + + slow_shared_task = make_task([group_a, group_b], slow_shared) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group_a, group_b], tasks=[failing_task, slow_shared_task]) + ) + + assert initial_groups == [group_a, group_b] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(group_a, boom), + GroupValuesEvent(group_b, ["shared-slow"]), + GroupSuccessEvent(group_b, [], []), + WorkQueueTerminationEvent(), + ] + + async def fails_all_pending_groups_when_a_shared_task_fails(): + group_a = Group(name="a") + group_b = Group(name="b") + + boom = RuntimeError("boom") + + def fail_a() -> WorkResult: + raise boom + + failing_task_a = make_task([group_a], fail_a) + + async def slow_fail_shared() -> WorkResult: + await sleep(0) + raise boom + + slow_failing_shared_task = make_task([group_a, group_b], slow_fail_shared) + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[group_a, group_b], + tasks=[failing_task_a, slow_failing_shared_task], + ) + ) + + assert initial_groups == [group_a, group_b] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(group_a, boom), + GroupFailureEvent(group_b, boom), + WorkQueueTerminationEvent(), + ] + + async def removes_failed_subtrees_with_already_removed_children(): + parent = Group(name="parent") + child_a = Group(parent, name="child_a") + child_b = Group(parent, name="child_b") + + boom = RuntimeError("boom") + + def fail() -> WorkResult: + raise boom + + failing_shared_task = make_task([child_a, parent], fail) + child_b_ran_spy = Spy(lambda: WorkResult("child_b")) + child_b_task = make_task([child_b], child_b_ran_spy) + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[parent, child_a, child_b], + tasks=[failing_shared_task, child_b_task], + ) + ) + + assert initial_groups == [parent] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(child_a, boom), + GroupFailureEvent(parent, boom), + WorkQueueTerminationEvent(), + ] + assert child_b_ran_spy.call_count == 0 + + async def integrates_late_work_referencing_dead_groups_gracefully(): + dead1 = Group(name="dead1") + dead2 = Group(name="dead2") + alive = Group(name="alive") + + boom1 = RuntimeError("boom1") + boom2 = RuntimeError("boom2") + + def fail1() -> WorkResult: + raise boom1 + + def fail2() -> WorkResult: + raise boom2 + + run_spy = Spy(lambda: WorkResult("never")) + late_stream = stream_from([WorkResult(1)], name="late") + + async def slow_dead() -> WorkResult: + await sleep(0) + await sleep(0) + return WorkResult( + "late", + Work( + groups=[Group(dead1, name="child_of_dead")], + tasks=[make_task([dead1, dead2], run_spy)], + streams=[late_stream], + ), + ) + + async def slow_alive() -> WorkResult: + await sleep(0) + await sleep(0) + await sleep(0) + await sleep(0) + return WorkResult("alive") + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[dead1, dead2, alive], + tasks=[ + make_task([dead1], fail1), + make_task([dead2], fail2), + make_task([dead1], slow_dead), + make_task([alive], slow_alive), + ], + ) + ) + + assert initial_groups == [dead1, dead2, alive] + assert initial_streams == [] + assert events == [ + GroupFailureEvent(dead1, boom1), + GroupFailureEvent(dead2, boom2), + GroupValuesEvent(alive, ["alive"]), + GroupSuccessEvent(alive, [], []), + WorkQueueTerminationEvent(), + ] + assert run_spy.call_count == 0 + + async def starts_tasks_added_to_a_pending_root_group_immediately(): + group = Group(name="group") + + async def slow() -> WorkResult: + await sleep(0) + await sleep(0) + return WorkResult("slow") + + extra_task = make_task([group], "extra") + stream = stream_from([WorkResult(1, Work(tasks=[extra_task]))]) + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[make_task([group], slow)], streams=[stream]) + ) + + assert initial_groups == [group] + assert initial_streams == [stream] + assert events == [ + StreamValuesEvent(stream, [1], [], []), + StreamSuccessEvent(stream), + GroupValuesEvent(group, ["slow", "extra"]), + GroupSuccessEvent(group, [], []), + WorkQueueTerminationEvent(), + ] + + async def prunes_duplicate_empty_child_group_entries(): + group = Group(name="group") + child = Group(group, name="child") + + async def slow_two() -> WorkResult: + await sleep(0) + return WorkResult("two", Work(groups=[child])) + + initial_groups, initial_streams, events = await collect_work_run( + Work( + groups=[group], + tasks=[ + make_task([group], "one", Work(groups=[child])), + make_task([group], slow_two), + ], + ) + ) + + assert initial_groups == [group] + assert initial_streams == [] + assert events == [ + GroupValuesEvent(group, ["one", "two"]), + GroupSuccessEvent(group, [], []), + WorkQueueTerminationEvent(), + ] + + async def merges_stream_success_when_the_queue_is_known_to_have_stopped(): + group = Group(name="group") + stream = Stream(FakeStreamQueue([[WorkResult(1)]], stops_early=True)) + + async def slow() -> WorkResult: + await sleep(0) + await sleep(0) + return WorkResult("slow") + + initial_groups, initial_streams, events = await collect_work_run( + Work(groups=[group], tasks=[make_task([group], slow)], streams=[stream]) + ) + + assert initial_groups == [group] + assert initial_streams == [stream] + # the stream success is delivered together with its values, + # and the redundant stream success signal from the pump is ignored + assert events == [ + StreamValuesEvent(stream, [1], [], []), + StreamSuccessEvent(stream), + GroupValuesEvent(group, ["slow"]), + GroupSuccessEvent(group, [], []), + WorkQueueTerminationEvent(), + ] + + def can_print_test_groups_and_streams(): + group = Group(name="some group") + assert repr(group) == "" + stream = Stream(FakeStreamQueue(), name="some stream") + assert repr(stream) == "" diff --git a/tests/execution/test_customize.py b/tests/execution/test_customize.py index ff6f55f1..7ce4f081 100644 --- a/tests/execution/test_customize.py +++ b/tests/execution/test_customize.py @@ -2,7 +2,7 @@ import pytest -from graphql.execution import Executor, GraphQLWrappedResult, execute, subscribe +from graphql.execution import Executor, execute, subscribe from graphql.language import parse from graphql.type import GraphQLField, GraphQLObjectType, GraphQLSchema, GraphQLString @@ -47,20 +47,17 @@ def execute_field( source, field_details_list, path, - incremental_context, - defer_map, + position_context, ): result = super().execute_field( parent_type, source, field_details_list, path, - incremental_context, - defer_map, + position_context, ) - assert isinstance(result, GraphQLWrappedResult) - result.result *= 2 - return result + assert isinstance(result, str) + return result * 2 assert execute( schema, diff --git a/tests/execution/test_executor.py b/tests/execution/test_executor.py index bbf71a0a..7f05d10d 100644 --- a/tests/execution/test_executor.py +++ b/tests/execution/test_executor.py @@ -17,7 +17,7 @@ from graphql.execution.get_variable_signature import GraphQLVariableSignature from graphql.execution.values import VariableValues, VariableValueSource from graphql.language import FieldNode, OperationDefinitionNode, parse -from graphql.pyutils import Undefined, inspect +from graphql.pyutils import Undefined, inspect, is_awaitable from graphql.type import ( GraphQLArgument, GraphQLBoolean, @@ -1366,3 +1366,68 @@ def memoizes_collect_subfields_results(): ) assert third is not first + + +def describe_base_executor_without_incremental_delivery(): + def deduplicates_errors_at_nulled_positions(): + from graphql.execution.executor import CollectedErrors + from graphql.pyutils import Path + + collected_errors = CollectedErrors() + path = Path(None, "foo", "Foo") + error = GraphQLError("first") + collected_errors.add(error, path) + assert collected_errors.errors == [error] + assert collected_errors.has_nulled_position(path) + # errors at or under an already nulled position are not added + collected_errors.add(GraphQLError("same position"), path) + collected_errors.add(GraphQLError("under position"), path.add_key("bar")) + assert collected_errors.errors == [error] + assert not collected_errors.has_nulled_position(Path(None, "other", "Foo")) + + def ignores_stream_directives(): + schema = GraphQLSchema( + GraphQLObjectType( + "Query", {"list": GraphQLField(GraphQLList(GraphQLString))} + ), + directives=[GraphQLStreamDirective], + ) + document = parse("{ list @stream(initialCount: 1) }") + executor = Executor.build(schema, document, {"list": ["a", "b", "c"]}) + assert isinstance(executor, Executor) + # the base executor completes the whole list, ignoring @stream + assert executor.execute_operation() == ({"list": ["a", "b", "c"]}, None) + + def ignores_defer_directives_on_subfields(): + obj_type = GraphQLObjectType("Obj", {"echo": GraphQLField(GraphQLString)}) + schema = GraphQLSchema( + GraphQLObjectType("Query", {"obj": GraphQLField(obj_type)}), + directives=[GraphQLDeferDirective], + ) + document = parse("{ obj { ... @defer { echo } } }") + executor = Executor.build(schema, document, {"obj": {"echo": "hello"}}) + assert isinstance(executor, Executor) + # the base executor executes deferred subfields inline + assert executor.execute_operation() == ({"obj": {"echo": "hello"}}, None) + + async def completes_async_is_type_of_with_async_subfields(): + async def is_obj(_value, _info): + return True + + async def echo(_value, _info): + return "hello" + + obj_type = GraphQLObjectType( + "Obj", + {"echo": GraphQLField(GraphQLString, resolve=echo)}, + is_type_of=is_obj, + ) + schema = GraphQLSchema( + GraphQLObjectType("Query", {"obj": GraphQLField(obj_type)}) + ) + document = parse("{ obj { echo } }") + executor = Executor.build(schema, document, {"obj": {}}) + assert isinstance(executor, Executor) + result = executor.execute_operation() + assert is_awaitable(result) + assert await result == ({"obj": {"echo": "hello"}}, None) diff --git a/tests/execution/test_executor_throwing_on_incremental.py b/tests/execution/test_executor_throwing_on_incremental.py new file mode 100644 index 00000000..a1e826b1 --- /dev/null +++ b/tests/execution/test_executor_throwing_on_incremental.py @@ -0,0 +1,120 @@ +"""Tests for the executor throwing on incremental delivery. + +The executor throwing on incremental delivery is used to execute operations +that must not produce multiple payloads. It is used for subscription events +by default; these tests exercise it directly for other operation types. +""" + +from __future__ import annotations + +from graphql import ( + GraphQLDeferDirective, + GraphQLField, + GraphQLList, + GraphQLObjectType, + GraphQLSchema, + GraphQLStreamDirective, + GraphQLString, + parse, + specified_directives, +) +from graphql.execution.executor_throwing_on_incremental import ( + ExecutorThrowingOnIncremental, +) + +obj_type = GraphQLObjectType( + "Obj", + { + "echo": GraphQLField(GraphQLString), + "list": GraphQLField(GraphQLList(GraphQLString)), + }, +) + +schema = GraphQLSchema( + GraphQLObjectType( + "Query", + { + "echo": GraphQLField(GraphQLString), + "obj": GraphQLField(obj_type), + "list": GraphQLField(GraphQLList(GraphQLString)), + }, + ), + subscription=GraphQLObjectType( + "Subscription", {"echo": GraphQLField(GraphQLString)} + ), + directives=[ + *specified_directives, + GraphQLDeferDirective, + GraphQLStreamDirective, + ], +) + +MULTIPLE_PAYLOADS_MESSAGE = ( + "Executing this GraphQL operation would unexpectedly produce" + " multiple payloads (due to @defer or @stream directive)" +) + + +def execute_throwing(source: str, root_value: dict | None = None): + executor = ExecutorThrowingOnIncremental.build(schema, parse(source), root_value) + assert isinstance(executor, ExecutorThrowingOnIncremental) + return executor.execute_operation() + + +def describe_executor_throwing_on_incremental(): + def executes_operations_without_incremental_delivery(): + result = execute_throwing("{ echo }", {"echo": "hello"}) + assert result == ({"echo": "hello"}, None) + + def completes_list_values_without_stream_directives(): + result = execute_throwing("{ list }", {"list": ["a", "b"]}) + assert result == ({"list": ["a", "b"]}, None) + + def raises_when_subscription_root_fields_are_deferred(): + result = execute_throwing( + "subscription { ... @defer { echo } }", {"echo": "hello"} + ) + assert result == ( + None, + [ + { + "message": "`@defer` directive not supported on subscription" + " operations. Disable `@defer` by setting the `if` argument" + " to `false`.", + } + ], + ) + + def raises_when_root_fields_are_deferred(): + result = execute_throwing("{ ... @defer { echo } }", {"echo": "hello"}) + assert result == (None, [{"message": MULTIPLE_PAYLOADS_MESSAGE}]) + + def raises_when_subfields_are_deferred(): + result = execute_throwing( + "{ obj { ... @defer { echo } } }", {"obj": {"echo": "hello"}} + ) + assert result == ( + {"obj": None}, + [ + { + "message": MULTIPLE_PAYLOADS_MESSAGE, + "locations": [(1, 3)], + "path": ["obj"], + } + ], + ) + + def raises_when_list_fields_are_streamed(): + result = execute_throwing( + "{ list @stream(initialCount: 1) }", {"list": ["a", "b"]} + ) + assert result == ( + {"list": None}, + [ + { + "message": MULTIPLE_PAYLOADS_MESSAGE, + "locations": [(1, 3)], + "path": ["list"], + } + ], + ) diff --git a/tests/pyutils/test_boxed_awaitabe_or_value.py b/tests/pyutils/test_boxed_awaitabe_or_value.py deleted file mode 100644 index 1b9f7474..00000000 --- a/tests/pyutils/test_boxed_awaitabe_or_value.py +++ /dev/null @@ -1,60 +0,0 @@ -from asyncio import Future, get_running_loop, isfuture, sleep - -import pytest - -from graphql.pyutils import BoxedAwaitableOrValue - -pytestmark = pytest.mark.anyio - - -async def async_function() -> int: - """A simple async function for testing awaitables.""" - return 42 - - -def create_future() -> Future: - """Create a new Future object.""" - return get_running_loop().create_future() - - -def describe_boxed_awaitable_futureor_value(): - def can_box_a_value(): - value = [42] - boxed: BoxedAwaitableOrValue[list[int]] = BoxedAwaitableOrValue(value) - - assert boxed.value is value - - async def can_box_a_future(): - future: Future[int] = create_future() - boxed: BoxedAwaitableOrValue[int] = BoxedAwaitableOrValue(future) - - assert boxed.value is future - - async def can_box_a_coroutine(): - awaitable = async_function() - - boxed: BoxedAwaitableOrValue[int] = BoxedAwaitableOrValue(awaitable) - assert isfuture(boxed.value) - - async def updates_value_when_future_is_done(): - future: Future[list[int]] = create_future() - boxed: BoxedAwaitableOrValue[list[int]] = BoxedAwaitableOrValue(future) - assert boxed.value is future - - value = [42] - future.set_result(value) - # value is updated immediately - assert boxed.value is value - - await sleep(0) - # value is still available - assert boxed.value is value - - async def updates_value_when_coroutine_was_awaited(): - awaitable = async_function() - - boxed: BoxedAwaitableOrValue[int] = BoxedAwaitableOrValue(awaitable) - assert isfuture(boxed.value) - - await sleep(0) # Allow the event loop to run the coroutine - assert boxed.value == 42