Fixed subscriptions if document and not query is provided

Author: syrusakbary

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

@@ -216,13 +216,20 @@ async def handle_subscribe(self, message: SubscribeMessage) -> None:
             await self.websocket.close(code=4401, reason="Unauthorized")
             return
 
-        try:
-            graphql_document = parse(message["payload"]["query"])
-        except GraphQLSyntaxError as exc:
-            await self.websocket.close(code=4400, reason=exc.message)
-            return
+        request_data = await self.view.get_graphql_request_data(
+            self.websocket, self.context, message["payload"], "subscription"
+        )
 
-        operation_name = message["payload"].get("operationName")
+        if request_data.document is not None:
+            graphql_document = request_data.document
+        else:
+            try:
+                graphql_document = parse(request_data.query)
+            except GraphQLSyntaxError as exc:
+                await self.websocket.close(code=4400, reason=exc.message)
+                return
+
+        operation_name = request_data.operation_name
 
         try:
             operation_type = get_operation_type(graphql_document, operation_name)
@@ -248,15 +255,11 @@ async def handle_subscribe(self, message: SubscribeMessage) -> None:
 
         if self.debug:  # pragma: no cover
             pretty_print_graphql_operation(
-                message["payload"].get("operationName"),
-                message["payload"]["query"],
-                message["payload"].get("variables"),
+                request_data.operation_name,
+                request_data.query,
+                request_data.variables,
             )
 
-        request_data = await self.view.get_graphql_request_data(
-            self.websocket, self.context, message["payload"], "subscription"
-        )
-
         operation = Operation(
             self,
             message["id"],