Make subscriptions use as well the document node

Author: syrusakbary

src/graphql_server/http/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from graphql.language import DocumentNode
 from typing import TYPE_CHECKING, Any, Optional
 from typing_extensions import Literal, TypedDict
 
@@ -30,10 +31,11 @@ class GraphQLRequestData:
     # query is optional here as it can be added by an extensions
     # (for example an extension for persisted queries)
     query: Optional[str]
+    document: Optional[DocumentNode]
     variables: Optional[dict[str, Any]]
     operation_name: Optional[str]
     extensions: Optional[dict[str, Any]]
-    protocol: Literal["http", "multipart-subscription"] = "http"
+    protocol: Literal["http", "multipart-subscription", "subscription"] = "http"
 
 
 __all__ = [

src/graphql_server/http/async_base_view.py

@@ -218,7 +218,7 @@ async def execute_operation(
         if request_data.protocol == "multipart-subscription":
             return await subscribe(
                 schema=self.schema,
-                query=request_data.query,  # type: ignore
+                query=request_data.document or request_data.query,  # type: ignore
                 variable_values=request_data.variables,
                 context_value=context,
                 root_value=root_value,
@@ -228,7 +228,7 @@ async def execute_operation(
 
         return await execute(
             schema=self.schema,
-            query=request_data.query,
+            query=request_data.document or request_data.query,
             root_value=root_value,
             variable_values=request_data.variables,
             context_value=context,
@@ -539,10 +539,13 @@ async def parse_multipart_subscriptions(
         return self.parse_json(await request.get_body())
 
     async def get_graphql_request_data(
-        self, data: dict[str, Any], protocol: Literal["http", "multipart-subscription"]
+        self,
+        data: dict[str, Any],
+        protocol: Literal["http", "multipart-subscription", "subscription"],
     ) -> GraphQLRequestData:
         return GraphQLRequestData(
             query=data.get("query"),
+            document=None,
             variables=data.get("variables"),
             operation_name=data.get("operationName"),
             extensions=data.get("extensions"),

src/graphql_server/http/sync_base_view.py

@@ -116,7 +116,7 @@ def execute_operation(
 
         return execute_sync(
             schema=self.schema,
-            query=request_data.query,
+            query=request_data.document or request_data.query,
             root_value=root_value,
             variable_values=request_data.variables,
             context_value=context,
@@ -139,6 +139,7 @@ def get_graphql_request_data(
     ) -> GraphQLRequestData:
         return GraphQLRequestData(
             query=data.get("query"),
+            document=None,
             variables=data.get("variables"),
             operation_name=data.get("operationName"),
             extensions=data.get("extensions"),

src/graphql_server/subscriptions/protocols/graphql_transport_ws/handlers.py

@@ -18,6 +18,7 @@
 
 from graphql_server import execute, subscribe
 from graphql_server.exceptions import ConnectionRejectionError, GraphQLValidationError
+from graphql_server.http import GraphQLRequestData
 from graphql_server.http.exceptions import (
     NonJsonMessageReceived,
     NonTextMessageReceived,
@@ -253,13 +254,15 @@ async def handle_subscribe(self, message: SubscribeMessage) -> None:
                 message["payload"].get("variables"),
             )
 
+        request_data = await self.view.get_graphql_request_data(
+            message["payload"], "subscription"
+        )
+
         operation = Operation(
             self,
             message["id"],
             operation_type,
-            message["payload"]["query"],
-            message["payload"].get("variables"),
-            message["payload"].get("operationName"),
+            request_data,
         )
 
         operation.task = asyncio.create_task(self.run_operation(operation))
@@ -274,7 +277,8 @@ async def run_operation(self, operation: Operation[Context, RootValue]) -> None:
             if operation.operation_type == OperationType.SUBSCRIPTION:
                 result_source = await subscribe(
                     schema=self.schema,
-                    query=operation.query,
+                    query=operation.request_data.document
+                    or operation.request_data.query,
                     variable_values=operation.variables,
                     operation_name=operation.operation_name,
                     context_value=self.context,
@@ -284,7 +288,8 @@ async def run_operation(self, operation: Operation[Context, RootValue]) -> None:
             else:
                 result_source = await execute(
                     schema=self.schema,
-                    query=operation.query,
+                    query=operation.request_data.document
+                    or operation.request_data.query,
                     variable_values=operation.variables,
                     context_value=self.context,
                     root_value=self.root_value,
@@ -378,31 +383,37 @@ class Operation(Generic[Context, RootValue]):
         "completed",
         "handler",
         "id",
-        "operation_name",
         "operation_type",
-        "query",
+        "request_data",
         "task",
-        "variables",
     ]
 
     def __init__(
         self,
         handler: BaseGraphQLTransportWSHandler[Context, RootValue],
         id: str,
         operation_type: OperationType,
-        query: str,
-        variables: Optional[dict[str, object]],
-        operation_name: Optional[str],
+        request_data: GraphQLRequestData,
     ) -> None:
         self.handler = handler
         self.id = id
         self.operation_type = operation_type
-        self.query = query
-        self.variables = variables
-        self.operation_name = operation_name
+        self.request_data = request_data
         self.completed = False
         self.task: Optional[asyncio.Task] = None
 
+    @property
+    def query(self) -> Optional[str]:
+        return self.request_data.query
+
+    @property
+    def variables(self) -> Optional[dict[str, Any]]:
+        return self.request_data.variables
+
+    @property
+    def operation_name(self) -> Optional[str]:
+        return self.request_data.operation_name
+
     async def send_operation_message(self, message: Message) -> None:
         if self.completed:
             return