Support schema extensions for subsctriptions - Alternative

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.types import Info
 from strawberry.extensions import SchemaExtension
@@ -48,6 +46,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

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
@@ -62,7 +72,7 @@ async def subscribe(
         context_value: Optional[Any] = None,
         root_value: Optional[Any] = None,
         operation_name: Optional[str] = None,
-    ) -> Any:
+    ) -> Union[ExecutionResult, AsyncGenerator[ExecutionResult, None]]:
         raise NotImplementedError
 
     @abstractmethod