Support subscriptions extensions

RELEASE.md

@@ -0,0 +1,116 @@
+Release type: minor
+
+Support for schema-extensions in subscriptions.
+
+i.e:
+```python
+import asyncio
+from typing import AsyncIterator
+
+import strawberry
+from strawberry.extensions.base_extension import SchemaExtension
+from strawberry.types.execution import PreExecutionError
+
+
+@strawberry.type
+class Query:
+    @strawberry.field
+    def hello(self) -> str:
+        return "Hello, world!"
+
+
+@strawberry.type
+class Subscription:
+    @strawberry.subscription
+    async def notifications(self, info: strawberry.Info) -> AsyncIterator[str]:
+        for _ in range(3):
+            yield "!dlrow ,olleH"
+
+
+class MyExtension(SchemaExtension):
+    def __init__(self):
+        self.count = 0
+
+    async def on_operation(self):
+        # This would run when the subscription starts
+        print("Subscription started")
+        yield
+        # The subscription has ended
+        print("Subscription ended")
+
+    async def on_execute(self):
+        # The subscription is trying to yield a new result
+        print(f"before yield {self.count}")
+        yield
+        # the subscription has yielded a new result
+        print(f"after yield {self.count}")
+        self.count += 1
+
+    # Other hooks are the same as in normal execution.
+    async def resolve(self, _next, root, info, *args, **kwargs):
+        res = _next(root, info, *args, **kwargs)
+        return res[::-1]
+
+
+schema = strawberry.Schema(
+    query=Query, subscription=Subscription, extensions=[MyExtension]
+)
+
+
+async def main():
+    agen = await schema.subscribe("subscription { notifications }")
+    if isinstance(agen, PreExecutionError):
+        print("this is an initial execution error.")
+        print(agen.errors)
+    else:
+        async for res in agen:
+            print(res.data)
+
+
+asyncio.run(main())
+```
+
+Should output this
+```console
+Subscription started
+before yield 0
+after yield 0
+before yield 1
+after yield 1
+{'notifications': 'Hello, world!'}
+before yield 2
+after yield 2
+{'notifications': 'Hello, world!'}
+before yield 3
+after yield 3
+{'notifications': 'Hello, world!'}
+before yield 4
+after yield 4
+Subscription ended
+```
+### Breaking changes
+This release also updates the signature of `Schema.subscribe`.
+From:
+```py
+async def subscribe(
+    self,
+    # TODO: make this optional when we support extensions
+    query: 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]:
+```
+To:
+```py
+async def subscribe(
+    self,
+    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[AsyncGenerator[ExecutionResult, None], PreExecutionError]:
+```
+Due to moving away from graphql-core result types to our internal types.

docs/guides/custom-extensions.mdx

@@ -150,7 +150,8 @@ class MyExtension(SchemaExtension):
 
 - Execution
 
-`on_execute` can be used to run code on the execution step of the GraphQL
+`on_execute` can be used to run code on the execution step of the GraphQL. When
+using subscriptions, `on_execute` will be called for each subscription yield.
 execution.
 
 ```python

noxfile.py

@@ -47,9 +47,10 @@ def _install_gql_core(session: Session, version: str) -> None:
     # hack for better workflow names  # noqa: FIX004
     if version == "3.2.3":
         session._session.install(f"graphql-core=={version}")  # type: ignore
-    session._session.install(
-        "https://github.com/graphql-python/graphql-core/archive/876aef67b6f1e1f21b3b5db94c7ff03726cb6bdf.zip"
-    )  # type: ignore
+    else:
+        session._session.install(
+            "https://github.com/graphql-python/graphql-core/archive/876aef67b6f1e1f21b3b5db94c7ff03726cb6bdf.zip"
+        )  # type: ignore
 
 
 gql_core_parametrize = nox.parametrize(

strawberry/channels/testing.py

@@ -144,9 +144,9 @@ async def subscribe(
             message_type = response["type"]
             if message_type == NextMessage.type:
                 payload = NextMessage(**response).payload
-                ret = ExecutionResult(payload["data"], None)
+                ret = ExecutionResult(payload.get("data"), None)
                 if "errors" in payload:
-                    ret.errors = self.process_errors(payload["errors"])
+                    ret.errors = self.process_errors(payload.get("errors") or [])
                 ret.extensions = payload.get("extensions", None)
                 yield ret
             elif message_type == ErrorMessage.type:

strawberry/extensions/base_extension.py

@@ -21,9 +21,9 @@ class LifecycleStep(Enum):
 class SchemaExtension:
     execution_context: ExecutionContext
 
-    def __init__(self, *, execution_context: ExecutionContext) -> None:
-        self.execution_context = execution_context
-
+    # to support extensions that still use the old signature
+    # we have an optional argument here for ease of initialization.
+    def __init__(self, execution_context: ExecutionContext | None = None) -> None: ...
     def on_operation(  # type: ignore
         self,
     ) -> AsyncIteratorOrIterator[None]:  # pragma: no cover

strawberry/extensions/context.py

@@ -24,8 +24,6 @@
 from strawberry.utils.await_maybe import AwaitableOrValue, await_maybe
 
 if TYPE_CHECKING:
-    from types import TracebackType
-
     from strawberry.extensions.base_extension import Hook
 
 
@@ -181,7 +179,7 @@ def __exit__(
         self,
         exc_type: Optional[Type[BaseException]],
         exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_tb: Optional[types.TracebackType],
     ) -> None:
         self.exit_stack.__exit__(exc_type, exc_val, exc_tb)
 
@@ -200,7 +198,7 @@ async def __aexit__(
         self,
         exc_type: Optional[Type[BaseException]],
         exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_tb: Optional[types.TracebackType],
     ) -> None:
         await self.async_exit_stack.__aexit__(exc_type, exc_val, exc_tb)
 

strawberry/extensions/runner.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
-
-from graphql import MiddlewareManager
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from strawberry.extensions.context import (
     ExecutingContextManager,
@@ -13,39 +11,22 @@
 )
 from strawberry.utils.await_maybe import await_maybe
 
-from . import SchemaExtension
-
 if TYPE_CHECKING:
     from strawberry.types import ExecutionContext
 
+    from . import SchemaExtension
+
 
 class SchemaExtensionsRunner:
     extensions: List[SchemaExtension]
 
     def __init__(
         self,
         execution_context: ExecutionContext,
-        extensions: Optional[
-            List[Union[Type[SchemaExtension], SchemaExtension]]
-        ] = None,
+        extensions: Optional[List[SchemaExtension]] = None,
     ) -> None:
         self.execution_context = execution_context
-
-        if not extensions:
-            extensions = []
-
-        init_extensions: List[SchemaExtension] = []
-
-        for extension in extensions:
-            # If the extension has already been instantiated then set the
-            # `execution_context` attribute
-            if isinstance(extension, SchemaExtension):
-                extension.execution_context = execution_context
-                init_extensions.append(extension)
-            else:
-                init_extensions.append(extension(execution_context=execution_context))
-
-        self.extensions = init_extensions
+        self.extensions = extensions or []
 
     def operation(self) -> OperationContextManager:
         return OperationContextManager(self.extensions)
@@ -61,29 +42,19 @@ def executing(self) -> ExecutingContextManager:
 
     def get_extensions_results_sync(self) -> Dict[str, Any]:
         data: Dict[str, Any] = {}
-
         for extension in self.extensions:
             if inspect.iscoroutinefunction(extension.get_results):
                 msg = "Cannot use async extension hook during sync execution"
                 raise RuntimeError(msg)
-
             data.update(extension.get_results())  # type: ignore
 
         return data
 
-    async def get_extensions_results(self) -> Dict[str, Any]:
+    async def get_extensions_results(self, ctx: ExecutionContext) -> Dict[str, Any]:
         data: Dict[str, Any] = {}
 
         for extension in self.extensions:
-            results = await await_maybe(extension.get_results())
-            data.update(results)
+            data.update(await await_maybe(extension.get_results()))
 
+        data.update(ctx.extensions_results)
         return data
-
-    def as_middleware_manager(self, *additional_middlewares: Any) -> MiddlewareManager:
-        middlewares = tuple(self.extensions) + additional_middlewares
-
-        return MiddlewareManager(*middlewares)
-
-
-__all__ = ["SchemaExtensionsRunner"]

strawberry/schema/base.py

@@ -20,6 +20,7 @@
     from strawberry.types.union import StrawberryUnion
 
     from .config import StrawberryConfig
+    from .subscribe import SubscriptionResult
 
 
 class BaseSchema(Protocol):
@@ -62,7 +63,7 @@ async def subscribe(
         context_value: Optional[Any] = None,
         root_value: Optional[Any] = None,
         operation_name: Optional[str] = None,
-    ) -> Any:
+    ) -> SubscriptionResult:
         raise NotImplementedError
 
     @abstractmethod