Support custom execution contexts in subscriptions (#181)

Author: bkad

src/graphql/execution/subscribe.py

@@ -1,5 +1,5 @@
 from inspect import isawaitable
-from typing import Any, AsyncIterable, AsyncIterator, Dict, Optional, Union
+from typing import Any, AsyncIterable, AsyncIterator, Dict, Optional, Type, Union
 
 from ..error import GraphQLError, located_error
 from ..execution.collect_fields import collect_fields
@@ -29,6 +29,7 @@ async def subscribe(
     operation_name: Optional[str] = None,
     field_resolver: Optional[GraphQLFieldResolver] = None,
     subscribe_field_resolver: Optional[GraphQLFieldResolver] = None,
+    execution_context_class: Optional[Type["ExecutionContext"]] = None,
 ) -> Union[AsyncIterator[ExecutionResult], ExecutionResult]:
     """Create a GraphQL subscription.
 
@@ -57,6 +58,7 @@ async def subscribe(
         variable_values,
         operation_name,
         subscribe_field_resolver,
+        execution_context_class,
     )
     if isinstance(result_or_stream, ExecutionResult):
         return result_or_stream
@@ -79,6 +81,7 @@ async def map_source_to_response(payload: Any) -> ExecutionResult:
             variable_values,
             operation_name,
             field_resolver,
+            execution_context_class=execution_context_class,
         )
         return await result if isawaitable(result) else result
 
@@ -94,6 +97,7 @@ async def create_source_event_stream(
     variable_values: Optional[Dict[str, Any]] = None,
     operation_name: Optional[str] = None,
     subscribe_field_resolver: Optional[GraphQLFieldResolver] = None,
+    execution_context_class: Optional[Type["ExecutionContext"]] = None,
 ) -> Union[AsyncIterable[Any], ExecutionResult]:
     """Create source event stream
 
@@ -122,9 +126,12 @@ async def create_source_event_stream(
     # mistake which should throw an early error.
     assert_valid_execution_arguments(schema, document, variable_values)
 
+    if not execution_context_class:
+        execution_context_class = ExecutionContext
+
     # If a valid context cannot be created due to incorrect arguments,
     # a "Response" with only errors is returned.
-    context = ExecutionContext.build(
+    context = execution_context_class.build(
         schema,
         document,
         root_value,

tests/execution/test_subscribe.py

@@ -3,7 +3,12 @@
 
 from pytest import mark, raises
 
-from graphql.execution import MapAsyncIterator, create_source_event_stream, subscribe
+from graphql.execution import (
+    create_source_event_stream,
+    subscribe,
+    MapAsyncIterator,
+    ExecutionContext,
+)
 from graphql.language import parse
 from graphql.pyutils import SimplePubSub
 from graphql.type import (
@@ -892,3 +897,42 @@ def resolve_message(message, _info):
         assert isinstance(subscription, MapAsyncIterator)
 
         assert await anext(subscription) == ({"newMessage": "Hello"}, None)
+
+    @mark.asyncio
+    async def should_work_with_custom_execution_contexts():
+        class CustomExecutionContext(ExecutionContext):
+            def build_resolve_info(self, *args, **kwargs):
+                resolve_info = super().build_resolve_info(*args, **kwargs)
+                resolve_info.context['foo'] = 'bar'
+                return resolve_info
+
+        async def generate_messages(_obj, info):
+            yield info.context['foo']
+
+        def resolve_message(message, _info):
+            return message
+
+        schema = GraphQLSchema(
+            query=QueryType,
+            subscription=GraphQLObjectType(
+                "Subscription",
+                {
+                    "newMessage": GraphQLField(
+                        GraphQLString,
+                        resolve=resolve_message,
+                        subscribe=generate_messages,
+                    )
+                },
+            ),
+        )
+
+        document = parse("subscription { newMessage }")
+        subscription = await subscribe(
+            schema,
+            document,
+            context_value={},
+            execution_context_class=CustomExecutionContext
+        )
+        assert isinstance(subscription, MapAsyncIterator)
+
+        assert await anext(subscription) == ({"newMessage": "bar"}, None)