Support subscriptions extensions

RELEASE.md

@@ -0,0 +1,36 @@
+Release type: minor
+
+This release adds support for schema-extensions in subscriptions.
+
+Here's a small example of how to use them (they work the same way as query and
+mutation extensions):
+
+```python
+import asyncio
+from typing import AsyncIterator
+
+import strawberry
+from strawberry.extensions.base_extension import SchemaExtension
+
+
+@strawberry.type
+class Subscription:
+ @strawberry.subscription
+ async def notifications(self, info: strawberry.Info) -> AsyncIterator[str]:
+ for _ in range(3):
+ yield "Hello"
+
+
+class MyExtension(SchemaExtension):
+ async def on_operation(self):
+ # This would run when the subscription starts
+ print("Subscription started")
+ yield
+ # The subscription has ended
+ print("Subscription ended")
+
+
+schema = strawberry.Schema(
+ query=Query, subscription=Subscription, extensions=[MyExtension]
+)
+```

TWEET.md

@@ -0,0 +1,5 @@
+🆕 Release $version is out! Thanks to $contributor for the PR 👏
+
+This release adds supports for schema extensions to subscriptions!
+
+Get it here 👉 $release_url

docs/breaking-changes.md

@@ -4,6 +4,7 @@ title: List of breaking changes and deprecations
 
 # List of breaking changes and deprecations
 
+- [Version 0.240.0 - 10 September 2024](./breaking-changes/0.240.0.md)
 - [Version 0.236.0 - 17 July 2024](./breaking-changes/0.236.0.md)
 - [Version 0.233.0 - 29 May 2024](./breaking-changes/0.233.0.md)
 - [Version 0.217.0 - 18 December 2023](./breaking-changes/0.217.0.md)

docs/breaking-changes/0.240.0.md

@@ -0,0 +1,36 @@
+---
+title: 0.240.0 Breaking Changes
+slug: breaking-changes/0.240.0
+---
+
+# v0.240.0 updates `Schema.subscribe`'s signature
+
+In order to support schema extensions in subscriptions and errors that can be
+raised before the execution of the subscription, we had to update the signature
+of `Schema.subscribe`.
+
+Previously it was:
+
+```python
+async 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,
+) -> Union[AsyncIterator[GraphQLExecutionResult], GraphQLExecutionResult]:
+```
+
+Now it is:
+
+```python
+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]:
+```

noxfile.py

@@ -11,7 +11,7 @@
 PYTHON_VERSIONS = ["3.12", "3.11", "3.10", "3.9", "3.8"]
 GQL_CORE_VERSIONS = [
  "3.2.3",
- "3.3.0",
+ "3.3.0a6",
 ]
 
 COMMON_PYTEST_OPTIONS = [
@@ -44,12 +44,7 @@
 
 
 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
+ session._session.install(f"graphql-core=={version}")
 
 
 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,11 @@ 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
@@ -61,6 +63,11 @@ def resolve(
  def get_results(self) -> AwaitableOrValue[Dict[str, Any]]:
  return {}
 
+ @classmethod
+ def _implements_resolve(cls) -> bool:
+ """Whether the extension implements the resolve method."""
+ return cls.resolve is not SchemaExtension.resolve
+
 
 Hook = Callable[[SchemaExtension], AsyncIteratorOrIterator[None]]
 

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/http/__init__.py

@@ -3,7 +3,7 @@
 import json
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional
-from typing_extensions import TypedDict
+from typing_extensions import Literal, TypedDict
 
 if TYPE_CHECKING:
  from strawberry.types import ExecutionResult
@@ -33,6 +33,7 @@ class GraphQLRequestData:
  query: Optional[str]
  variables: Optional[Dict[str, Any]]
  operation_name: Optional[str]
+ protocol: Literal["http", "multipart-subscription"] = "http"
 
 
 def parse_query_params(params: Dict[str, str]) -> Dict[str, Any]:

strawberry/http/async_base_view.py

@@ -14,6 +14,7 @@
  Tuple,
  Union,
 )
+from typing_extensions import Literal
 
 from graphql import GraphQLError
 
@@ -121,6 +122,15 @@ async def execute_operation(
 
  assert self.schema
 
+ if request_data.protocol == "multipart-subscription":
+ return await self.schema.subscribe(
+ request_data.query, # type: ignore
+ variable_values=request_data.variables,
+ context_value=context,
+ root_value=root_value,
+ operation_name=request_data.operation_name,
+ )
+
  return await self.schema.execute(
  request_data.query,
  root_value=root_value,
@@ -312,21 +322,27 @@ async def parse_http_body(
  ) -> GraphQLRequestData:
  content_type, params = parse_content_type(request.content_type or "")
 
+ protocol: Literal["http", "multipart-subscription"] = "http"
+
  if request.method == "GET":
  data = self.parse_query_params(request.query_params)
+ if self._is_multipart_subscriptions(content_type, params):
+ protocol = "multipart-subscription"
  elif "application/json" in content_type:
  data = self.parse_json(await request.get_body())
  elif content_type == "multipart/form-data":
  data = await self.parse_multipart(request)
  elif self._is_multipart_subscriptions(content_type, params):
  data = await self.parse_multipart_subscriptions(request)
+ protocol = "multipart-subscription"
  else:
  raise HTTPException(400, "Unsupported content type")
 
  return GraphQLRequestData(
  query=data.get("query"),
  variables=data.get("variables"),
  operation_name=data.get("operationName"),
+ protocol=protocol,
  )
 
  async def process_result(

strawberry/schema/base.py

@@ -15,7 +15,6 @@
  from strawberry.types import (
  ExecutionContext,
  ExecutionResult,
- SubscriptionExecutionResult,
  )
  from strawberry.types.base import StrawberryObjectDefinition
  from strawberry.types.enum import EnumDefinition
@@ -24,6 +23,7 @@
  from strawberry.types.union import StrawberryUnion
 
  from .config import StrawberryConfig
+ from .subscribe import SubscriptionResult
 
 
 class BaseSchema(Protocol):
@@ -43,7 +43,7 @@ async def execute(
  root_value: Optional[Any] = None,
  operation_name: Optional[str] = None,
  allowed_operation_types: Optional[Iterable[OperationType]] = None,
- ) -> Union[ExecutionResult, SubscriptionExecutionResult]:
+ ) -> ExecutionResult:
  raise NotImplementedError
 
  @abstractmethod
@@ -66,7 +66,7 @@ async def subscribe(
  context_value: Optional[Any] = None,
  root_value: Optional[Any] = None,
  operation_name: Optional[str] = None,
- ) -> Any:
+ ) -> SubscriptionResult:
  raise NotImplementedError
 
  @abstractmethod