Support Schema Extensions in subscriptions

RELEASE.md

@@ -0,0 +1,3 @@
+Release type: minor
+
+Subscriptions now support Schema Extensions.

docs/guides/custom-extensions.md

@@ -36,8 +36,6 @@ resolvers.
 If you need to wrap only certain field resolvers with additional logic, please
 check out [field extensions](field-extensions.md).
 
-Note that `resolve` can also be implemented asynchronously.
-
 ```python
 from strawberry.extensions import SchemaExtension
 
@@ -47,6 +45,21 @@ class MyExtension(SchemaExtension):
         return _next(root, info, *args, **kwargs)
 ```
 
+Note that `resolve` can also be implemented asynchronously, in which case the
+result from `_next` must be optionally awaited:
+
+```python
+from inspect import isawaitable
+from strawberry.types import Info
+from strawberry.extensions import SchemaExtension
+
+
+class MyExtension(SchemaExtension):
+    async def resolve(self, _next, root, info: Info, *args, **kwargs):
+        result = _next(root, info, *args, **kwargs)
+        return await result if isawaitable(result) else result
+```
+
 ### Get results
 
 `get_results` allows to return a dictionary of data or alternatively an

docs/operations/testing.md

@@ -136,10 +136,8 @@ async def test_subscription():
     	}
     """
 
-    sub = await schema.subscribe(query)
-
     index = 0
-    async for result in sub:
+    async for result in schema.subscribe(query):
         assert not result.errors
         assert result.data == {"count": index}
 

strawberry/schema/__init__.py

@@ -1,4 +1,4 @@
-from .base import BaseSchema
+from .base import BaseSchema, SubscribeSingleResult
 from .schema import Schema
 
-__all__ = ["BaseSchema", "Schema"]
+__all__ = ["BaseSchema", "Schema", "SubscribeSingleResult"]

strawberry/schema/base.py

@@ -2,7 +2,17 @@
 
 from abc import abstractmethod
 from functools import lru_cache
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Type,
+    Union,
+)
 from typing_extensions import Protocol
 
 from strawberry.utils.logging import StrawberryLogger
@@ -22,6 +32,15 @@
     from .config import StrawberryConfig
 
 
+class SubscribeSingleResult(RuntimeError):
+    """Raised when Schema.subscribe() returns a single execution result, instead of a
+    subscription generator, typically as a result of validation errors.
+    """
+
+    def __init__(self, value: ExecutionResult) -> None:
+        self.value = value
+
+
 class BaseSchema(Protocol):
     config: StrawberryConfig
     schema_converter: GraphQLCoreConverter
@@ -55,14 +74,14 @@ def execute_sync(
         raise NotImplementedError
 
     @abstractmethod
-    async def subscribe(
+    def subscribe(
         self,
         query: str,
         variable_values: Optional[Dict[str, Any]] = None,
         context_value: Optional[Any] = None,
         root_value: Optional[Any] = None,
         operation_name: Optional[str] = None,
-    ) -> Any:
+    ) -> AsyncGenerator[ExecutionResult, None]:
         raise NotImplementedError
 
     @abstractmethod

strawberry/schema/execute.py

@@ -4,6 +4,8 @@
 from inspect import isawaitable
 from typing import (
     TYPE_CHECKING,
+    AsyncGenerator,
+    AsyncIterable,
     Awaitable,
     Callable,
     Iterable,
@@ -17,21 +19,23 @@
     cast,
 )
 
+from graphql import ExecutionResult as GraphQLExecutionResult
 from graphql import GraphQLError, parse
 from graphql import execute as original_execute
+from graphql import subscribe as original_subscribe
 from graphql.validation import validate
 
 from strawberry.exceptions import MissingQueryError
 from strawberry.extensions.runner import SchemaExtensionsRunner
 from strawberry.types import ExecutionResult
 
+from .base import SubscribeSingleResult
 from .exceptions import InvalidOperationTypeError
 
 if TYPE_CHECKING:
     from typing_extensions import NotRequired, Unpack
 
     from graphql import ExecutionContext as GraphQLExecutionContext
-    from graphql import ExecutionResult as GraphQLExecutionResult
     from graphql import GraphQLSchema
     from graphql.language import DocumentNode
     from graphql.validation import ASTValidationRule
@@ -272,3 +276,149 @@ def execute_sync(
         errors=execution_context.result.errors,
         extensions=extensions_runner.get_extensions_results_sync(),
     )
+
+
+async def subscribe(
+    schema: GraphQLSchema,
+    *,
+    extensions: Sequence[Union[Type[SchemaExtension], SchemaExtension]],
+    execution_context: ExecutionContext,
+    process_errors: Callable[[List[GraphQLError], Optional[ExecutionContext]], None],
+) -> AsyncGenerator[ExecutionResult, None]:
+    """
+    The graphql-core subscribe function returns either an ExecutionResult or an
+    AsyncGenerator[ExecutionResult, None].  The former is returned in case of an error
+    during parsing or validation.
+    Because we need to maintain execution context, we cannot return an
+    async generator, we must _be_ an async generator.  So we yield a
+    (bool, ExecutionResult) tuple, where the bool indicates whether the result is an
+    potentially multiple execution result or a single result.
+    A False value indicates an single result, most likely an intial
+    failure (and no more values will be yielded) whereas a True value indicates a
+    successful subscription, and more values may be yielded.
+    """
+
+    extensions_runner = SchemaExtensionsRunner(
+        execution_context=execution_context,
+        extensions=list(extensions),
+    )
+
+    # unlike execute(), the entire operation, including the results hooks,
+    # is run within the operation() hook.
+    async with extensions_runner.operation():
+        # Note: In graphql-core the schema would be validated here but in
+        # Strawberry we are validating it at initialisation time instead
+        assert execution_context.query is not None
+
+        async with extensions_runner.parsing():
+            try:
+                if not execution_context.graphql_document:
+                    execution_context.graphql_document = parse_document(
+                        execution_context.query, **execution_context.parse_options
+                    )
+
+            except GraphQLError as error:
+                execution_context.errors = [error]
+                process_errors([error], execution_context)
+                raise SubscribeSingleResult(
+                    ExecutionResult(
+                        data=None,
+                        errors=[error],
+                        extensions=await extensions_runner.get_extensions_results(),
+                    )
+                )
+
+            except Exception as error:  # pragma: no cover
+                error = GraphQLError(str(error), original_error=error)
+                execution_context.errors = [error]
+                process_errors([error], execution_context)
+
+                raise SubscribeSingleResult(
+                    ExecutionResult(
+                        data=None,
+                        errors=[error],
+                        extensions=await extensions_runner.get_extensions_results(),
+                    )
+                )
+
+        async with extensions_runner.validation():
+            _run_validation(execution_context)
+            if execution_context.errors:
+                process_errors(execution_context.errors, execution_context)
+                raise SubscribeSingleResult(
+                    ExecutionResult(data=None, errors=execution_context.errors)
+                )
+
+        async with extensions_runner.executing():
+            # currently original_subscribe is an async function.  A future release
+            # of graphql-core will make it optionally awaitable
+            result: Union[AsyncIterable[GraphQLExecutionResult], GraphQLExecutionResult]
+            result_or_awaitable = original_subscribe(
+                schema,
+                execution_context.graphql_document,
+                root_value=execution_context.root_value,
+                context_value=execution_context.context,
+                variable_values=execution_context.variables,
+                operation_name=execution_context.operation_name,
+            )
+            if isawaitable(result_or_awaitable):
+                result = await cast(
+                    Awaitable[
+                        Union[
+                            AsyncIterable["GraphQLExecutionResult"],
+                            "GraphQLExecutionResult",
+                        ]
+                    ],
+                    result_or_awaitable,
+                )
+            else:  # pragma: no cover
+                result = cast(
+                    Union[
+                        AsyncIterable["GraphQLExecutionResult"],
+                        "GraphQLExecutionResult",
+                    ],
+                    result_or_awaitable,
+                )
+
+            if isinstance(result, GraphQLExecutionResult):
+                raise SubscribeSingleResult(
+                    await process_subscribe_result(
+                        execution_context, process_errors, extensions_runner, result
+                    )
+                )
+
+            aiterator = result.__aiter__()
+            try:
+                async for result in aiterator:
+                    yield await process_subscribe_result(
+                        execution_context, process_errors, extensions_runner, result
+                    )
+            finally:
+                # grapql-core's iterator may or may not have an aclose() method
+                if hasattr(aiterator, "aclose"):
+                    await aiterator.aclose()
+
+
+async def process_subscribe_result(
+    execution_context: ExecutionContext,
+    process_errors: Callable[[List[GraphQLError], Optional[ExecutionContext]], None],
+    extensions_runner: SchemaExtensionsRunner,
+    result: GraphQLExecutionResult,
+) -> ExecutionResult:
+    execution_context.result = result
+    # Also set errors on the execution_context so that it's easier
+    # to access in extensions
+    if result.errors:
+        execution_context.errors = result.errors
+
+        # Run the `Schema.process_errors` function here before
+        # extensions have a chance to modify them (see the MaskErrors
+        # extension). That way we can log the original errors but
+        # only return a sanitised version to the client.
+        process_errors(result.errors, execution_context)
+
+    return ExecutionResult(
+        data=execution_context.result.data,
+        errors=execution_context.result.errors,
+        extensions=await extensions_runner.get_extensions_results(),
+    )

strawberry/schema/schema.py

@@ -5,7 +5,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncIterator,
+    AsyncGenerator,
     Dict,
     Iterable,
     List,
@@ -20,10 +20,8 @@
     GraphQLNonNull,
     GraphQLSchema,
     get_introspection_query,
-    parse,
     validate_schema,
 )
-from graphql.execution import subscribe
 from graphql.type.directives import specified_directives
 
 from strawberry import relay
@@ -43,11 +41,10 @@
 from . import compat
 from .base import BaseSchema
 from .config import StrawberryConfig
-from .execute import execute, execute_sync
+from .execute import execute, execute_sync, subscribe
 
 if TYPE_CHECKING:
     from graphql import ExecutionContext as GraphQLExecutionContext
-    from graphql import ExecutionResult as GraphQLExecutionResult
 
     from strawberry.custom_scalar import ScalarDefinition, ScalarWrapper
     from strawberry.directive import StrawberryDirective
@@ -305,22 +302,29 @@ def execute_sync(
 
     async def subscribe(
         self,
-        # TODO: make this optional when we support extensions
-        query: str,
+        query: Optional[str],
         variable_values: Optional[Dict[str, Any]] = None,
         context_value: Optional[Any] = None,
         root_value: Optional[Any] = None,
         operation_name: Optional[str] = None,
-    ) -> Union[AsyncIterator[GraphQLExecutionResult], GraphQLExecutionResult]:
-        return await subscribe(
-            self._schema,
-            parse(query),
+    ) -> AsyncGenerator[ExecutionResult, None]:
+        execution_context = ExecutionContext(
+            query=query,
+            schema=self,
+            context=context_value,
             root_value=root_value,
-            context_value=context_value,
-            variable_values=variable_values,
-            operation_name=operation_name,
+            variables=variable_values,
+            provided_operation_name=operation_name,
         )
 
+        async for result in subscribe(
+            self._schema,
+            extensions=self.get_extensions(),
+            execution_context=execution_context,
+            process_errors=self.process_errors,
+        ):
+            yield result
+
     def _resolve_node_ids(self):
         for concrete_type in self.schema_converter.type_map.values():
             type_def = concrete_type.definition