feat: extend subscription with middleware

Author: KarolJagodzinski

ariadne/asgi/handlers/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from inspect import isawaitable
 from logging import Logger, LoggerAdapter
-from typing import Any
+from typing import Any, cast
 
 from graphql import DocumentNode, ExecutionContext, GraphQLSchema, MiddlewareManager
 from starlette.types import Receive, Scope, Send
@@ -11,7 +11,11 @@
 from ...types import (
     ContextValue,
     ErrorFormatter,
+    ExtensionList,
+    Extensions,
     GraphQLResult,
+    MiddlewareList,
+    Middlewares,
     OnComplete,
     OnConnect,
     OnDisconnect,
@@ -168,6 +172,9 @@ def __init__(
         on_disconnect: OnDisconnect | None = None,
         on_operation: OnOperation | None = None,
         on_complete: OnComplete | None = None,
+        extensions: Extensions | None = None,
+        middleware: Middlewares | None = None,
+        middleware_manager_class: type[MiddlewareManager] | None = None,
     ) -> None:
         """Initialize websocket handler with optional options specific to it.
 
@@ -183,6 +190,19 @@ def __init__(
 
         `on_complete`: an `OnComplete` callback, used when GraphQL operation
         received over the websocket connection was completed.
+
+        `extensions`: an `Extensions` list or callable returning a
+        list of extensions server should use during subscription execution.
+        Defaults to no extensions.
+
+        `middleware`: a `Middlewares` list or callable returning a list of
+        middlewares server should use during subscription execution. Defaults
+        to no middlewares.
+
+        `middleware_manager_class`: a `MiddlewareManager` type or subclass to
+        use for combining provided middlewares into single wrapper for resolvers
+        by the server. Defaults to `graphql.MiddlewareManager`. Is only used
+        if `extensions` or `middleware` options are set.
         """
         super().__init__()
         self.http_handler: GraphQLHttpHandlerBase | None = None
@@ -191,11 +211,57 @@ def __init__(
         self.on_disconnect: OnDisconnect | None = on_disconnect
         self.on_operation: OnOperation | None = on_operation
         self.on_complete: OnComplete | None = on_complete
+        self.extensions = extensions
+        self.middleware = middleware
+        if middleware_manager_class is not None:
+            self.middleware_manager_class = middleware_manager_class
 
     @abstractmethod
     async def handle_websocket(self, websocket: Any):
         """Abstract method for handling the websocket connection."""
 
+    async def get_extensions_for_request(
+        self, request: Any, context: ContextValue | None
+    ) -> ExtensionList:
+        """Returns extensions to use when handling the GraphQL request.
+
+        Returns `ExtensionList`, a list of extensions to use or `None`.
+
+        # Required arguments
+
+        `request`: the `WebSocket` instance from Starlette or FastAPI.
+
+        `context`: a `ContextValue` for this request.
+        """
+        if callable(self.extensions):
+            extensions = self.extensions(request, context)  # ty: ignore
+            if isawaitable(extensions):
+                extensions = await extensions
+            return cast(ExtensionList, extensions)
+        return self.extensions
+
+    async def get_middleware_for_request(
+        self, request: Any, context: ContextValue | None
+    ) -> MiddlewareList:
+        """Returns GraphQL middlewares to use when handling the GraphQL request.
+
+        Returns `MiddlewareList`, a list of middlewares to use or `None`.
+
+        # Required arguments
+
+        `request`: the `WebSocket` instance from Starlette or FastAPI.
+
+        `context`: a `ContextValue` for this request.
+        """
+        middleware = self.middleware
+        if callable(middleware):
+            middleware = middleware(request, context)  # ty: ignore
+            if isawaitable(middleware):
+                middleware = await middleware
+        if middleware:
+            return cast(MiddlewareList, middleware)
+        return None
+
     def configure(
         self,
         *args,

ariadne/asgi/handlers/graphql_transport_ws.py

@@ -339,6 +339,9 @@ async def handle_websocket_subscribe(
                     "schema is not set, call configure method to initialize it"
                 )
 
+            extensions = await self.get_extensions_for_request(websocket, context_value)
+            middleware = await self.get_middleware_for_request(websocket, context_value)
+
             success, results_producer = await subscribe(
                 self.schema,
                 data,
@@ -351,6 +354,9 @@ async def handle_websocket_subscribe(
                 introspection=self.introspection,
                 logger=self.logger,
                 error_formatter=self.error_formatter,
+                extensions=extensions,
+                middleware=middleware,
+                middleware_manager_class=self.middleware_manager_class,
             )
         else:
             if self.http_handler is None:

ariadne/asgi/handlers/graphql_ws.py

@@ -296,6 +296,9 @@ async def start_websocket_operation(
         if self.schema is None:
             raise TypeError("schema is not set, call configure method to initialize it")
 
+        extensions = await self.get_extensions_for_request(websocket, context_value)
+        middleware = await self.get_middleware_for_request(websocket, context_value)
+
         success, results = await subscribe(
             self.schema,
             data,
@@ -308,6 +311,9 @@ async def start_websocket_operation(
             introspection=self.introspection,
             logger=self.logger,
             error_formatter=self.error_formatter,
+            extensions=extensions,
+            middleware=middleware,
+            middleware_manager_class=self.middleware_manager_class,
         )
 
         if not success:

ariadne/asgi/handlers/http.py

@@ -205,6 +205,12 @@ async def graphql_http_server(self, request: Request) -> Response:
             for handler in self.subscription_handlers:
                 if handler.supports(request, data):
                     context_value = await self.get_context_for_request(request, data)
+                    extensions = await self.get_extensions_for_request(
+                        request, context_value
+                    )
+                    middleware = await self.get_middleware_for_request(
+                        request, context_value
+                    )
                     return await handler.handle(
                         request,
                         data,
@@ -218,6 +224,9 @@ async def graphql_http_server(self, request: Request) -> Response:
                         introspection=self.introspection,
                         logger=self.logger,
                         error_formatter=self.error_formatter,
+                        extensions=extensions,
+                        middleware=middleware,
+                        middleware_manager_class=self.middleware_manager_class,
                     )
 
         success, result = await self.execute_graphql_query(request, data)

ariadne/contrib/federation/utils.py

@@ -1,3 +1,4 @@
+import copy
 from inspect import isawaitable
 from typing import Any, cast
 
@@ -49,29 +50,40 @@ def _purge_directive_nodes(nodes: tuple[Node, ...]) -> tuple[Node, ...]:
     )
 
 
-def _purge_type_directives(definition: Node):
-    # Recursively check every field defined on the Node definition
-    # and remove any directives found.
+def _purge_type_directives(definition: Node) -> Node:
+    """Recursively check every field defined on the Node definition
+    and remove any directives found. Returns a new node with purged directives."""
+    changes: dict[str, Any] = {}
     for key in definition.keys:
         value = getattr(definition, key, None)
         if isinstance(value, tuple):
             # Remove directive nodes from the tuple
             # e.g. doc -> definitions [DirectiveDefinitionNode]
             next_value = _purge_directive_nodes(cast(tuple[Node, ...], value))
-            for item in next_value:
-                if isinstance(item, Node):
-                    # Look for directive nodes on sub-nodes, e.g.: doc ->
-                    # definitions [ObjectTypeDefinitionNode] -> fields -> directives
-                    _purge_type_directives(item)
-            setattr(definition, key, next_value)
+            # Look for directive nodes on sub-nodes, e.g.: doc ->
+            # definitions [ObjectTypeDefinitionNode] -> fields -> directives
+            next_value = tuple(
+                _purge_type_directives(item) if isinstance(item, Node) else item
+                for item in next_value
+            )
+            if next_value != value:
+                changes[key] = next_value
         elif isinstance(value, Node):
-            _purge_type_directives(value)
+            new_value = _purge_type_directives(value)
+            if new_value is not value:
+                changes[key] = new_value
+    if changes:
+        new_node = copy.copy(definition)
+        for key, value in changes.items():
+            object.__setattr__(new_node, key, value)
+        return new_node
+    return definition
 
 
 def purge_schema_directives(joined_type_defs: str) -> str:
     """Remove custom schema directives from federation."""
     ast_document = parse(joined_type_defs)
-    _purge_type_directives(ast_document)
+    ast_document = _purge_type_directives(ast_document)
     return print_ast(ast_document)
 
 

ariadne/contrib/sse.py

@@ -41,7 +41,9 @@
 from ..subscription_handlers.handlers import SubscriptionHandler
 from ..types import (
     ErrorFormatter,
+    ExtensionList,
     Extensions,
+    MiddlewareList,
     Middlewares,
     QueryParser,
     QueryValidator,
@@ -386,6 +388,9 @@ async def handle(
         introspection: bool,
         logger: None | str | Logger | LoggerAdapter,
         error_formatter: ErrorFormatter,
+        middleware: MiddlewareList = None,
+        middleware_manager_class: type[MiddlewareManager] | None = None,
+        extensions: ExtensionList | None = None,
     ) -> Response:
         """Handle the subscription request via Server-Sent Events.
 
@@ -405,6 +410,9 @@ async def handle(
                 introspection=introspection,
                 logger=logger,
                 error_formatter=error_formatter,
+                middleware=middleware,
+                middleware_manager_class=middleware_manager_class,
+                extensions=extensions,
             ),
             ping_interval=self.ping_interval,
             send_timeout=self.send_timeout,

ariadne/contrib/tracing/utils.py

@@ -59,9 +59,10 @@ def repr_upload_file(upload_file: UploadFile | File) -> str:
 
 def format_path(path: ResponsePath):
     elements = []
-    while path:
-        elements.append(path.key)
-        path = path.prev
+    current: ResponsePath | None = path
+    while current:
+        elements.append(current.key)
+        current = current.prev
     return elements[::-1]
 
 

ariadne/graphql.py

@@ -441,6 +441,9 @@ async def subscribe(
     logger: None | str | Logger | LoggerAdapter = None,
     validation_rules: ValidationRules | None = None,
     error_formatter: ErrorFormatter = format_error,
+    middleware: MiddlewareList = None,
+    middleware_manager_class: type[MiddlewareManager] | None = None,
+    extensions: ExtensionList | None = None,
     **kwargs,
 ) -> SubscriptionResult:
     """Subscribe to GraphQL updates.
@@ -491,9 +494,20 @@ async def subscribe(
     `error_formatter`: an `ErrorFormatter` callable to use to convert GraphQL
     errors encountered during query execution to JSON-serializable format.
 
+    `middleware`: a `list` of or callable returning list of GraphQL middleware
+    to use by query executor.
+
+    `middleware_manager_class`: a `MiddlewareManager` class to use by query
+    executor.
+
+    `extensions`: a `list` of or callable returning list of extensions
+    to use during query execution.
+
     `**kwargs`: any kwargs not used by `subscribe` are passed to
     `graphql.subscribe`.
     """
+    extension_manager = ExtensionManager(extensions, context_value)
+
     try:
         validate_data(data)
         variables, operation_name = (
@@ -533,15 +547,21 @@ async def subscribe(
             if isawaitable(root_value):
                 root_value = await root_value
 
-        result = await _subscribe(
+        result = _subscribe(
             schema,
             document,
             root_value=root_value,
             context_value=context_value,
             variable_values=variables,
             operation_name=operation_name,
+            middleware=extension_manager.as_middleware_manager(
+                middleware, middleware_manager_class
+            ),
             **kwargs,
         )
+
+        if isawaitable(result):
+            result = await result
     except GraphQLError as error:
         log_error(error, logger)
         return False, [error_formatter(error, debug)]

ariadne/schema_visitor.py

@@ -725,9 +725,8 @@ def heal_type(
             # any `GraphQLNamedType` with a `name`, then it must end up identical
             # to `schema.get_type(name)`, since `schema.type_map` is the source
             # of truth for all named schema types.
-            named_type = cast(GraphQLNamedType, type_)
-            official_type = schema.get_type(named_type.name)
-            if official_type and named_type != official_type:
+            official_type = schema.get_type(type_.name)
+            if official_type and type_ != official_type:
                 return official_type
 
         return type_

ariadne/subscription_handlers/handlers.py

@@ -5,14 +5,22 @@
 from logging import Logger, LoggerAdapter
 from typing import Any, cast
 
-from graphql import DocumentNode, ExecutionResult, GraphQLError, GraphQLSchema
+from graphql import (
+    DocumentNode,
+    ExecutionResult,
+    GraphQLError,
+    GraphQLSchema,
+    MiddlewareManager,
+)
 from starlette.requests import Request
 from starlette.responses import Response
 
 from ..graphql import subscribe, validate_data
 from ..logger import log_error
 from ..types import (
     ErrorFormatter,
+    ExtensionList,
+    MiddlewareList,
     QueryParser,
     QueryValidator,
     RootValue,
@@ -59,6 +67,9 @@ async def handle(
         introspection: bool,
         logger: None | str | Logger | LoggerAdapter,
         error_formatter: ErrorFormatter,
+        middleware: MiddlewareList = None,
+        middleware_manager_class: type[MiddlewareManager] | None = None,
+        extensions: ExtensionList | None = None,
     ) -> Response:
         """Handle the subscription request."""
 
@@ -77,6 +88,9 @@ async def generate_events(
         introspection: bool,
         logger: None | str | Logger | LoggerAdapter,
         error_formatter: ErrorFormatter,
+        middleware: MiddlewareList = None,
+        middleware_manager_class: type[MiddlewareManager] | None = None,
+        extensions: ExtensionList | None = None,
     ) -> AsyncGenerator[SubscriptionEvent, None]:
         """Execute subscription and yield events.
 
@@ -105,6 +119,9 @@ async def generate_events(
                 introspection=introspection,
                 logger=logger,
                 error_formatter=error_formatter,
+                middleware=middleware,
+                middleware_manager_class=middleware_manager_class,
+                extensions=extensions,
             )
 
             if not success: