strawberry-graphql 0.239.1__py3-none-any.whl → 0.240.0__py3-none-any.whl

strawberry/schema/schema.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
 import warnings
-from functools import lru_cache
+from functools import cached_property, lru_cache
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncIterator,
     Dict,
     Iterable,
     List,
@@ -22,18 +21,19 @@ from graphql import (
     GraphQLNonNull,
     GraphQLSchema,
     get_introspection_query,
-    parse,
     validate_schema,
 )
-from graphql.execution import subscribe
+from graphql.execution.middleware import MiddlewareManager
 from graphql.type.directives import specified_directives
 
 from strawberry import relay
 from strawberry.annotation import StrawberryAnnotation
+from strawberry.extensions import SchemaExtension
 from strawberry.extensions.directives import (
     DirectivesExtension,
     DirectivesExtensionSync,
 )
+from strawberry.extensions.runner import SchemaExtensionsRunner
 from strawberry.schema.schema_converter import GraphQLCoreConverter
 from strawberry.schema.types.scalar import DEFAULT_SCALAR_REGISTRY
 from strawberry.types import ExecutionContext
@@ -41,19 +41,17 @@ from strawberry.types.base import StrawberryObjectDefinition, has_object_definit
 from strawberry.types.graphql import OperationType
 
 from ..printer import print_schema
-from ..utils.await_maybe import await_maybe
 from . import compat
 from .base import BaseSchema
 from .config import StrawberryConfig
 from .execute import execute, execute_sync
+from .subscribe import SubscriptionResult, subscribe
 
 if TYPE_CHECKING:
     from graphql import ExecutionContext as GraphQLExecutionContext
-    from graphql import ExecutionResult as GraphQLExecutionResult
 
     from strawberry.directive import StrawberryDirective
-    from strawberry.extensions import SchemaExtension
-    from strawberry.types import ExecutionResult, SubscriptionExecutionResult
+    from strawberry.types import ExecutionResult
     from strawberry.types.base import StrawberryType
     from strawberry.types.enum import EnumDefinition
     from strawberry.types.field import StrawberryField
@@ -85,7 +83,7 @@ class Schema(BaseSchema):
         ] = None,
         schema_directives: Iterable[object] = (),
     ) -> None:
-        """Default Schema to be to be used in a Strawberry application.
+        """Default Schema to be used in a Strawberry application.
 
         A GraphQL Schema class used to define the structure and configuration
         of GraphQL queries, mutations, and subscriptions.
@@ -125,6 +123,7 @@ class Schema(BaseSchema):
         self.subscription = subscription
 
         self.extensions = extensions
+        self._cached_middleware_manager: MiddlewareManager | None = None
         self.execution_context_class = execution_context_class
         self.config = config or StrawberryConfig()
 
@@ -214,15 +213,63 @@ class Schema(BaseSchema):
             formatted_errors = "\n\n".join(f"❌ {error.message}" for error in errors)
             raise ValueError(f"Invalid Schema. Errors:\n\n{formatted_errors}")
 
-    def get_extensions(
-        self, sync: bool = False
-    ) -> List[Union[Type[SchemaExtension], SchemaExtension]]:
-        extensions = list(self.extensions)
-
+    def get_extensions(self, sync: bool = False) -> List[SchemaExtension]:
+        extensions = []
         if self.directives:
-            extensions.append(DirectivesExtensionSync if sync else DirectivesExtension)
+            extensions = [
+                *self.extensions,
+                DirectivesExtensionSync if sync else DirectivesExtension,
+            ]
+        extensions.extend(self.extensions)
+        return [
+            ext if isinstance(ext, SchemaExtension) else ext(execution_context=None)
+            for ext in extensions
+        ]
+
+    @cached_property
+    def _sync_extensions(self) -> List[SchemaExtension]:
+        return self.get_extensions(sync=True)
+
+    @cached_property
+    def _async_extensions(self) -> List[SchemaExtension]:
+        return self.get_extensions(sync=False)
+
+    def create_extensions_runner(
+        self, execution_context: ExecutionContext, extensions: list[SchemaExtension]
+    ) -> SchemaExtensionsRunner:
+        return SchemaExtensionsRunner(
+            execution_context=execution_context,
+            extensions=extensions,
+        )
 
-        return extensions
+    def _get_middleware_manager(
+        self, extensions: list[SchemaExtension]
+    ) -> MiddlewareManager:
+        # create a middleware manager with all the extensions that implement resolve
+        if not self._cached_middleware_manager:
+            self._cached_middleware_manager = MiddlewareManager(
+                *(ext for ext in extensions if ext._implements_resolve())
+            )
+        return self._cached_middleware_manager
+
+    def _create_execution_context(
+        self,
+        query: Optional[str],
+        allowed_operation_types: Iterable[OperationType],
+        variable_values: Optional[Dict[str, Any]] = None,
+        context_value: Optional[Any] = None,
+        root_value: Optional[Any] = None,
+        operation_name: Optional[str] = None,
+    ) -> ExecutionContext:
+        return ExecutionContext(
+            query=query,
+            schema=self,
+            allowed_operations=allowed_operation_types,
+            context=context_value,
+            root_value=root_value,
+            variables=variable_values,
+            provided_operation_name=operation_name,
+        )
 
     @lru_cache
     def get_type_by_name(
@@ -284,31 +331,33 @@ class Schema(BaseSchema):
         root_value: Optional[Any] = None,
         operation_name: Optional[str] = None,
         allowed_operation_types: Optional[Iterable[OperationType]] = None,
-    ) -> Union[ExecutionResult, SubscriptionExecutionResult]:
+    ) -> ExecutionResult:
         if allowed_operation_types is None:
             allowed_operation_types = DEFAULT_ALLOWED_OPERATION_TYPES
 
-        # Create execution context
-        execution_context = ExecutionContext(
+        execution_context = self._create_execution_context(
             query=query,
-            schema=self,
-            context=context_value,
+            allowed_operation_types=allowed_operation_types,
+            variable_values=variable_values,
+            context_value=context_value,
             root_value=root_value,
-            variables=variable_values,
-            provided_operation_name=operation_name,
+            operation_name=operation_name,
         )
-
-        result = await execute(
+        extensions = self.get_extensions()
+        # TODO (#3571): remove this when we implement execution context as parameter.
+        for extension in extensions:
+            extension.execution_context = execution_context
+        return await execute(
             self._schema,
-            extensions=self.get_extensions(),
-            execution_context_class=self.execution_context_class,
             execution_context=execution_context,
-            allowed_operation_types=allowed_operation_types,
+            extensions_runner=self.create_extensions_runner(
+                execution_context, extensions
+            ),
             process_errors=self._process_errors,
+            middleware_manager=self._get_middleware_manager(extensions),
+            execution_context_class=self.execution_context_class,
         )
 
-        return result
-
     def execute_sync(
         self,
         query: Optional[str],
@@ -321,44 +370,59 @@ class Schema(BaseSchema):
         if allowed_operation_types is None:
             allowed_operation_types = DEFAULT_ALLOWED_OPERATION_TYPES
 
-        execution_context = ExecutionContext(
+        execution_context = self._create_execution_context(
             query=query,
-            schema=self,
-            context=context_value,
+            allowed_operation_types=allowed_operation_types,
+            variable_values=variable_values,
+            context_value=context_value,
             root_value=root_value,
-            variables=variable_values,
-            provided_operation_name=operation_name,
+            operation_name=operation_name,
         )
-
-        result = execute_sync(
+        extensions = self._sync_extensions
+        # TODO (#3571): remove this when we implement execution context as parameter.
+        for extension in extensions:
+            extension.execution_context = execution_context
+        return execute_sync(
             self._schema,
-            extensions=self.get_extensions(sync=True),
-            execution_context_class=self.execution_context_class,
             execution_context=execution_context,
+            extensions_runner=self.create_extensions_runner(
+                execution_context, extensions
+            ),
+            execution_context_class=self.execution_context_class,
             allowed_operation_types=allowed_operation_types,
             process_errors=self._process_errors,
+            middleware_manager=self._get_middleware_manager(extensions),
         )
 
-        return result
-
     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 await_maybe(
-            subscribe(
-                self._schema,
-                parse(query),
-                root_value=root_value,
-                context_value=context_value,
-                variable_values=variable_values,
-                operation_name=operation_name,
-            )
+    ) -> SubscriptionResult:
+        execution_context = self._create_execution_context(
+            query=query,
+            allowed_operation_types=(OperationType.SUBSCRIPTION,),
+            variable_values=variable_values,
+            context_value=context_value,
+            root_value=root_value,
+            operation_name=operation_name,
+        )
+        extensions = self._async_extensions
+        # TODO (#3571): remove this when we implement execution context as parameter.
+        for extension in extensions:
+            extension.execution_context = execution_context
+        return await subscribe(
+            self._schema,
+            execution_context=execution_context,
+            extensions_runner=self.create_extensions_runner(
+                execution_context, extensions
+            ),
+            process_errors=self._process_errors,
+            middleware_manager=self._get_middleware_manager(extensions),
+            execution_context_class=self.execution_context_class,
         )
 
     def _resolve_node_ids(self) -> None:

strawberry/schema/subscribe.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, AsyncGenerator, AsyncIterator, Optional, Type, Union
+
+from graphql import (
+    ExecutionResult as OriginalExecutionResult,
+)
+from graphql.execution import ExecutionContext as GraphQLExecutionContext
+from graphql.execution import subscribe as original_subscribe
+
+from strawberry.types import ExecutionResult
+from strawberry.types.execution import ExecutionContext, PreExecutionError
+from strawberry.utils import IS_GQL_32
+from strawberry.utils.await_maybe import await_maybe
+
+from .execute import (
+    ProcessErrors,
+    _coerce_error,
+    _handle_execution_result,
+    _parse_and_validate_async,
+)
+
+if TYPE_CHECKING:
+    from typing_extensions import TypeAlias
+
+    from graphql.execution.middleware import MiddlewareManager
+    from graphql.type.schema import GraphQLSchema
+
+    from ..extensions.runner import SchemaExtensionsRunner
+
+SubscriptionResult: TypeAlias = Union[
+    PreExecutionError, AsyncGenerator[ExecutionResult, None]
+]
+
+OriginSubscriptionResult = Union[
+    OriginalExecutionResult,
+    AsyncIterator[OriginalExecutionResult],
+]
+
+
+async def _subscribe(
+    schema: GraphQLSchema,
+    execution_context: ExecutionContext,
+    extensions_runner: SchemaExtensionsRunner,
+    process_errors: ProcessErrors,
+    middleware_manager: MiddlewareManager,
+    execution_context_class: Optional[Type[GraphQLExecutionContext]] = None,
+) -> AsyncGenerator[Union[PreExecutionError, ExecutionResult], None]:
+    async with extensions_runner.operation():
+        if initial_error := await _parse_and_validate_async(
+            context=execution_context,
+            extensions_runner=extensions_runner,
+        ):
+            initial_error.extensions = await extensions_runner.get_extensions_results(
+                execution_context
+            )
+            yield await _handle_execution_result(
+                execution_context, initial_error, extensions_runner, process_errors
+            )
+        try:
+            async with extensions_runner.executing():
+                assert execution_context.graphql_document is not None
+                gql_33_kwargs = {
+                    "middleware": middleware_manager,
+                    "execution_context_class": execution_context_class,
+                }
+                try:
+                    # Might not be awaitable for pre-execution errors.
+                    aiter_or_result: OriginSubscriptionResult = await await_maybe(
+                        original_subscribe(
+                            schema,
+                            execution_context.graphql_document,
+                            root_value=execution_context.root_value,
+                            variable_values=execution_context.variables,
+                            operation_name=execution_context.operation_name,
+                            context_value=execution_context.context,
+                            **{} if IS_GQL_32 else gql_33_kwargs,  # type: ignore[arg-type]
+                        )
+                    )
+                # graphql-core 3.2 doesn't handle some of the pre-execution errors.
+                # see `test_subscription_immediate_error`
+                except Exception as exc:
+                    aiter_or_result = OriginalExecutionResult(
+                        data=None, errors=[_coerce_error(exc)]
+                    )
+
+            # Handle pre-execution errors.
+            if isinstance(aiter_or_result, OriginalExecutionResult):
+                yield await _handle_execution_result(
+                    execution_context,
+                    PreExecutionError(data=None, errors=aiter_or_result.errors),
+                    extensions_runner,
+                    process_errors,
+                )
+            else:
+                try:
+                    async for result in aiter_or_result:
+                        yield await _handle_execution_result(
+                            execution_context,
+                            result,
+                            extensions_runner,
+                            process_errors,
+                        )
+                # graphql-core doesn't handle exceptions raised while executing.
+                except Exception as exc:
+                    yield await _handle_execution_result(
+                        execution_context,
+                        OriginalExecutionResult(data=None, errors=[_coerce_error(exc)]),
+                        extensions_runner,
+                        process_errors,
+                    )
+        # catch exceptions raised in `on_execute` hook.
+        except Exception as exc:
+            origin_result = OriginalExecutionResult(
+                data=None, errors=[_coerce_error(exc)]
+            )
+            yield await _handle_execution_result(
+                execution_context,
+                origin_result,
+                extensions_runner,
+                process_errors,
+            )
+
+
+async def subscribe(
+    schema: GraphQLSchema,
+    execution_context: ExecutionContext,
+    extensions_runner: SchemaExtensionsRunner,
+    process_errors: ProcessErrors,
+    middleware_manager: MiddlewareManager,
+    execution_context_class: Optional[Type[GraphQLExecutionContext]] = None,
+) -> SubscriptionResult:
+    asyncgen = _subscribe(
+        schema,
+        execution_context,
+        extensions_runner,
+        process_errors,
+        middleware_manager,
+        execution_context_class,
+    )
+    # GrapQL-core might return an initial error result instead of an async iterator.
+    # This happens when "there was an immediate error" i.e resolver is not an async iterator.
+    # To overcome this while maintaining the extension contexts we do this trick.
+    first = await asyncgen.__anext__()
+    if isinstance(first, PreExecutionError):
+        await asyncgen.aclose()
+        return first
+
+    async def _wrapper() -> AsyncGenerator[ExecutionResult, None]:
+        yield first
+        async for result in asyncgen:
+            yield result
+
+    return _wrapper()

strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py

@@ -7,15 +7,13 @@ from contextlib import suppress
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncGenerator,
-    AsyncIterator,
+    Awaitable,
     Callable,
     Dict,
     List,
     Optional,
 )
 
-from graphql import ExecutionResult as GraphQLExecutionResult
 from graphql import GraphQLError, GraphQLSyntaxError, parse
 
 from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
@@ -24,11 +22,14 @@ from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
     ConnectionInitMessage,
     ErrorMessage,
     NextMessage,
+    NextPayload,
     PingMessage,
     PongMessage,
     SubscribeMessage,
     SubscribeMessagePayload,
 )
+from strawberry.types import ExecutionResult
+from strawberry.types.execution import PreExecutionError
 from strawberry.types.graphql import OperationType
 from strawberry.types.unset import UNSET
 from strawberry.utils.debug import pretty_print_graphql_operation
@@ -38,10 +39,10 @@ if TYPE_CHECKING:
     from datetime import timedelta
 
     from strawberry.schema import BaseSchema
+    from strawberry.schema.subscribe import SubscriptionResult
     from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
         GraphQLTransportMessage,
     )
-    from strawberry.types import ExecutionResult
 
 
 class BaseGraphQLTransportWSHandler(ABC):
@@ -243,10 +244,10 @@ class BaseGraphQLTransportWSHandler(ABC):
         if isinstance(context, dict):
             context["connection_params"] = self.connection_params
         root_value = await self.get_root_value()
-
+        result_source: Awaitable[ExecutionResult] | Awaitable[SubscriptionResult]
         # Get an AsyncGenerator yielding the results
         if operation_type == OperationType.SUBSCRIPTION:
-            result_source = await self.schema.subscribe(
+            result_source = self.schema.subscribe(
                 query=message.payload.query,
                 variable_values=message.payload.variables,
                 operation_name=message.payload.operationName,
@@ -254,29 +255,16 @@ class BaseGraphQLTransportWSHandler(ABC):
                 root_value=root_value,
             )
         else:
-            # create AsyncGenerator returning a single result
-            async def get_result_source() -> AsyncIterator[ExecutionResult]:
-                yield await self.schema.execute(  # type: ignore
-                    query=message.payload.query,
-                    variable_values=message.payload.variables,
-                    context_value=context,
-                    root_value=root_value,
-                    operation_name=message.payload.operationName,
-                )
-
-            result_source = get_result_source()
+            result_source = self.schema.execute(
+                query=message.payload.query,
+                variable_values=message.payload.variables,
+                context_value=context,
+                root_value=root_value,
+                operation_name=message.payload.operationName,
+            )
 
         operation = Operation(self, message.id, operation_type)
 
-        # Handle initial validation errors
-        if isinstance(result_source, GraphQLExecutionResult):
-            assert operation_type == OperationType.SUBSCRIPTION
-            assert result_source.errors
-            payload = [err.formatted for err in result_source.errors]
-            await self.send_message(ErrorMessage(id=message.id, payload=payload))
-            self.schema.process_errors(result_source.errors)
-            return
-
         # Create task to handle this subscription, reserve the operation ID
         operation.task = asyncio.create_task(
             self.operation_task(result_source, operation)
@@ -284,65 +272,37 @@ class BaseGraphQLTransportWSHandler(ABC):
         self.operations[message.id] = operation
 
     async def operation_task(
-        self, result_source: AsyncGenerator, operation: Operation
+        self,
+        result_source: Awaitable[ExecutionResult] | Awaitable[SubscriptionResult],
+        operation: Operation,
     ) -> None:
         """The operation task's top level method. Cleans-up and de-registers the operation once it is done."""
         # TODO: Handle errors in this method using self.handle_task_exception()
         try:
-            await self.handle_async_results(result_source, operation)
-        except BaseException:  # pragma: no cover
-            # cleanup in case of something really unexpected
-            # wait for generator to be closed to ensure that any existing
-            # 'finally' statement is called
-            with suppress(RuntimeError):
-                await result_source.aclose()
-            if operation.id in self.operations:
-                del self.operations[operation.id]
+            first_res_or_agen = await result_source
+            # that's an immediate error we should end the operation
+            # without a COMPLETE message
+            if isinstance(first_res_or_agen, PreExecutionError):
+                assert first_res_or_agen.errors
+                await operation.send_initial_errors(first_res_or_agen.errors)
+            # that's a mutation / query result
+            elif isinstance(first_res_or_agen, ExecutionResult):
+                await operation.send_next(first_res_or_agen)
+                await operation.send_message(CompleteMessage(id=operation.id))
+            else:
+                async for result in first_res_or_agen:
+                    await operation.send_next(result)
+                await operation.send_message(CompleteMessage(id=operation.id))
+
+        except BaseException as e:  # pragma: no cover
+            self.operations.pop(operation.id, None)
             raise
-        else:
-            await operation.send_message(CompleteMessage(id=operation.id))
         finally:
             # add this task to a list to be reaped later
             task = asyncio.current_task()
             assert task is not None
             self.completed_tasks.append(task)
 
-    async def handle_async_results(
-        self,
-        result_source: AsyncGenerator,
-        operation: Operation,
-    ) -> None:
-        try:
-            async for result in result_source:
-                if (
-                    result.errors
-                    and operation.operation_type != OperationType.SUBSCRIPTION
-                ):
-                    error_payload = [err.formatted for err in result.errors]
-                    error_message = ErrorMessage(id=operation.id, payload=error_payload)
-                    await operation.send_message(error_message)
-                    # don't need to call schema.process_errors() here because
-                    # it was already done by schema.execute()
-                    return
-                else:
-                    next_payload = {"data": result.data}
-                    if result.errors:
-                        self.schema.process_errors(result.errors)
-                        next_payload["errors"] = [
-                            err.formatted for err in result.errors
-                        ]
-                    next_message = NextMessage(id=operation.id, payload=next_payload)
-                    await operation.send_message(next_message)
-        except Exception as error:
-            # GraphQLErrors are handled by graphql-core and included in the
-            # ExecutionResult
-            error = GraphQLError(str(error), original_error=error)
-            error_payload = [error.formatted]
-            error_message = ErrorMessage(id=operation.id, payload=error_payload)
-            await operation.send_message(error_message)
-            self.schema.process_errors([error])
-            return
-
     def forget_id(self, id: str) -> None:
         # de-register the operation id making it immediately available
         # for re-use
@@ -401,5 +361,21 @@ class Operation:
             self.handler.forget_id(self.id)
         await self.handler.send_message(message)
 
+    async def send_initial_errors(self, errors: list[GraphQLError]) -> None:
+        # Initial errors see https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md#error
+        # "This can occur before execution starts,
+        # usually due to validation errors, or during the execution of the request"
+        await self.send_message(
+            ErrorMessage(id=self.id, payload=[err.formatted for err in errors])
+        )
+
+    async def send_next(self, execution_result: ExecutionResult) -> None:
+        next_payload: NextPayload = {"data": execution_result.data}
+        if execution_result.errors:
+            next_payload["errors"] = [err.formatted for err in execution_result.errors]
+        if execution_result.extensions:
+            next_payload["extensions"] = execution_result.extensions
+        await self.send_message(NextMessage(id=self.id, payload=next_payload))
+
 
 __all__ = ["BaseGraphQLTransportWSHandler", "Operation"]

strawberry/subscriptions/protocols/graphql_transport_ws/types.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import asdict, dataclass
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict
 
 from strawberry.types.unset import UNSET
 
@@ -68,12 +68,20 @@ class SubscribeMessage(GraphQLTransportMessage):
     type: str = "subscribe"
 
 
+class NextPayload(TypedDict, total=False):
+    data: Any
+
+    # Optional list of formatted graphql.GraphQLError objects
+    errors: Optional[List[GraphQLFormattedError]]
+    extensions: Optional[Dict[str, Any]]
+
+
 @dataclass
 class NextMessage(GraphQLTransportMessage):
     """Direction: Server -> Client."""
 
     id: str
-    payload: Dict[str, Any]  # TODO: shape like FormattedExecutionResult
+    payload: NextPayload
     type: str = "next"
 
     def as_dict(self) -> dict: