Support subscriptions extensions

RELEASE.md

@@ -0,0 +1,56 @@
+Release type: minor
+
+Support for schema-extensions in subscriptions.
+
+i.e:
+```python
+from typing import AsyncIterator, List
+
+import strawberry
+from strawberry.extensions.base_extension import SchemaExtension
+
+
+@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]:
+ async for notification in info.context["pubsub"].subscribe(
+ "notifications", info.context["user"]
+ ):
+ yield notification
+
+
+class MyExtension(SchemaExtension):
+
+ async def on_operation(self):
+ # This would run when the subscription starts
+ self.monitor_start(self.execution_context.operation_name)
+ yield
+ # The subscription has ended
+ self.monitor_end(self.execution_context.operation_name)
+
+ async def on_execute(self):
+ count = 0
+ # The subsctription is trying to yield a new result
+ self.monitor_start(
+ self.execution_context.operation_name, f"before yield {count}"
+ )
+ yield
+ # the subscription has yielded a new result
+ self.monitor_end(self.execution_context.operation_name, f"after yield {count}")
+ count += 1
+
+ # other hooks can are the same as per normal execution.
+
+
+schema = strawberry.Schema(
+ query=Query, subscription=Subscription, extensions=[MyExtension()]
+)
+```

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

pyproject.toml

@@ -35,7 +35,8 @@ build-backend = "poetry.masonry.api"
 
 [tool.poetry.dependencies]
 python = "^3.8"
-graphql-core = "~3.2.0"
+# TOOD: revert this
+graphql-core = {git = "https://github.com/graphql-python/graphql-core.git", branch = "main"}
 typing-extensions = ">=4.5.0"
 python-dateutil = "^2.7.0"
 starlette = {version = ">=0.18.0", optional = true}

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/context.py

@@ -16,17 +16,15 @@
  List,
  NamedTuple,
  Optional,
- Type,
  Union,
 )
 
 from strawberry.extensions import SchemaExtension
 from strawberry.utils.await_maybe import AwaitableOrValue, await_maybe
 
 if TYPE_CHECKING:
- from types import TracebackType
-
  from strawberry.extensions.base_extension import Hook
+ from strawberry.types.execution import ExecutionContext
 
 
 class WrappedHook(NamedTuple):
@@ -56,7 +54,9 @@ def __init_subclass__(cls) -> None:
  LEGACY_ENTER: str
  LEGACY_EXIT: str
 
- def __init__(self, extensions: List[SchemaExtension]) -> None:
+ def __init__(
+ self, extensions: List[SchemaExtension], context_context: ExecutionContext
+ ) -> None:
  self.hooks: List[WrappedHook] = []
  self.default_hook: Hook = getattr(SchemaExtension, self.HOOK_NAME)
  for extension in extensions:
@@ -177,13 +177,8 @@ def __enter__(self) -> None:
  else:
  self.exit_stack.enter_context(hook.hook()) # type: ignore
 
- def __exit__(
- self,
- exc_type: Optional[Type[BaseException]],
- exc_val: Optional[BaseException],
- exc_tb: Optional[TracebackType],
- ) -> None:
- self.exit_stack.__exit__(exc_type, exc_val, exc_tb)
+ def __exit__(self, *_: object) -> None:
+ self.exit_stack.close()
 
  async def __aenter__(self) -> None:
  self.async_exit_stack = contextlib.AsyncExitStack()
@@ -196,13 +191,8 @@ async def __aenter__(self) -> None:
  else:
  self.async_exit_stack.enter_context(hook.hook()) # type: ignore
 
- async def __aexit__(
- self,
- exc_type: Optional[Type[BaseException]],
- exc_val: Optional[BaseException],
- exc_tb: Optional[TracebackType],
- ) -> None:
- await self.async_exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+ async def __aexit__(self, *_: object) -> None:
+ await self.async_exit_stack.aclose()
 
 
 class OperationContextManager(ExtensionContextManagerBase):

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,
@@ -25,65 +23,45 @@ class SchemaExtensionsRunner:
  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 if extensions is not None else []
 
  def operation(self) -> OperationContextManager:
- return OperationContextManager(self.extensions)
+ return OperationContextManager(self.extensions, self.execution_context)
 
  def validation(self) -> ValidationContextManager:
- return ValidationContextManager(self.extensions)
+ return ValidationContextManager(self.extensions, self.execution_context)
 
  def parsing(self) -> ParsingContextManager:
- return ParsingContextManager(self.extensions)
+ return ParsingContextManager(self.extensions, self.execution_context)
 
  def executing(self) -> ExecutingContextManager:
- return ExecutingContextManager(self.extensions)
+ return ExecutingContextManager(self.extensions, self.execution_context)
+
+ @classmethod
+ def _implments_get_rseults(cls, extension: SchemaExtension) -> bool:
+ """Whether the extension implements get_results."""
+ return type(extension).get_results is not SchemaExtension.get_results
 
  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
+ if type(self)._implments_get_rseults(extension):
+ 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)
+ if type(self)._implments_get_rseults(extension):
+  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"]