Support graphql-transport-ws websocket subprotocol

Author: SmileyChris

graphql_ws/base.py

@@ -4,13 +4,17 @@
 from graphql import format_error, graphql
 
 from .constants import (
+    GQL_COMPLETE,
     GQL_CONNECTION_ERROR,
     GQL_CONNECTION_INIT,
     GQL_CONNECTION_TERMINATE,
     GQL_DATA,
     GQL_ERROR,
+    GQL_NEXT,
     GQL_START,
     GQL_STOP,
+    GQL_SUBSCRIBE,
+    TRANSPORT_WS_PROTOCOL,
 )
 
 
@@ -23,6 +27,9 @@ def __init__(self, ws, request_context=None):
         self.ws = ws
         self.operations = {}
         self.request_context = request_context
+        self.transport_ws_protocol = request_context and TRANSPORT_WS_PROTOCOL in (
+            request_context.get("subprotocols") or []
+        )
 
     def has_operation(self, op_id):
         return op_id in self.operations
@@ -41,7 +48,7 @@ def remove_operation(self, op_id):
 
     def unsubscribe(self, op_id):
         async_iterator = self.remove_operation(op_id)
-        if hasattr(async_iterator, 'dispose'):
+        if hasattr(async_iterator, "dispose"):
             async_iterator.dispose()
         return async_iterator
 
@@ -84,12 +91,16 @@ def process_message(self, connection_context, parsed_message):
         elif op_type == GQL_CONNECTION_TERMINATE:
             return self.on_connection_terminate(connection_context, op_id)
 
-        elif op_type == GQL_START:
+        elif op_type == (
+            GQL_SUBSCRIBE if connection_context.transport_ws_protocol else GQL_START
+        ):
             assert isinstance(payload, dict), "The payload must be a dict"
             params = self.get_graphql_params(connection_context, payload)
             return self.on_start(connection_context, op_id, params)
 
-        elif op_type == GQL_STOP:
+        elif op_type == (
+            GQL_COMPLETE if connection_context.transport_ws_protocol else GQL_STOP
+        ):
             return self.on_stop(connection_context, op_id)
 
         else:
@@ -142,7 +153,12 @@ def build_message(self, id, op_type, payload):
 
     def send_execution_result(self, connection_context, op_id, execution_result):
         result = self.execution_result_to_dict(execution_result)
-        return self.send_message(connection_context, op_id, GQL_DATA, result)
+        return self.send_message(
+            connection_context,
+            op_id,
+            GQL_NEXT if connection_context.transport_ws_protocol else GQL_DATA,
+            result,
+        )
 
     def execution_result_to_dict(self, execution_result):
         result = OrderedDict()

graphql_ws/constants.py

@@ -1,5 +1,6 @@
 GRAPHQL_WS = "graphql-ws"
 WS_PROTOCOL = GRAPHQL_WS
+TRANSPORT_WS_PROTOCOL = "graphql-transport-ws"
 
 GQL_CONNECTION_INIT = "connection_init"  # Client -> Server
 GQL_CONNECTION_ACK = "connection_ack"  # Server -> Client
@@ -8,8 +9,10 @@
 # NOTE: This one here don't follow the standard due to connection optimization
 GQL_CONNECTION_TERMINATE = "connection_terminate"  # Client -> Server
 GQL_CONNECTION_KEEP_ALIVE = "ka"  # Server -> Client
-GQL_START = "start"  # Client -> Server
-GQL_DATA = "data"  # Server -> Client
+GQL_START = "start"  # Client -> Server  (graphql-ws)
+GQL_SUBSCRIBE = "subscribe"  # Client -> Server  (graphql-transport-ws START equivalent)
+GQL_DATA = "data"  # Server -> Client  (graphql-ws)
+GQL_NEXT = "next"  # Server -> Client  (graphql-transport-ws DATA equivalent)
 GQL_ERROR = "error"  # Server -> Client
-GQL_COMPLETE = "complete"  # Server -> Client
-GQL_STOP = "stop"  # Client -> Server
+GQL_COMPLETE = "complete"  # Server -> Client  (and Client -> Server for graphql-transport-ws STOP equivalent)
+GQL_STOP = "stop"  # Client -> Server  (graphql-ws only)

graphql_ws/django/consumers.py

@@ -2,21 +2,25 @@
 
 from channels.generic.websocket import AsyncJsonWebsocketConsumer
 
-from ..constants import WS_PROTOCOL
+from ..constants import TRANSPORT_WS_PROTOCOL, WS_PROTOCOL
 from .subscriptions import subscription_server
 
 
 class GraphQLSubscriptionConsumer(AsyncJsonWebsocketConsumer):
-
     async def connect(self):
         self.connection_context = None
-        if WS_PROTOCOL in self.scope["subprotocols"]:
-            self.connection_context = await subscription_server.handle(
-                ws=self, request_context=self.scope
-            )
-            await self.accept(subprotocol=WS_PROTOCOL)
-        else:
+        found_protocol = None
+        for protocol in [WS_PROTOCOL, TRANSPORT_WS_PROTOCOL]:
+            if protocol in self.scope["subprotocols"]:
+                found_protocol = protocol
+                break
+        if not found_protocol:
             await self.close()
+            return
+        self.connection_context = await subscription_server.handle(
+            ws=self, request_context=self.scope
+        )
+        await self.accept(subprotocol=found_protocol)
 
     async def disconnect(self, code):
         if self.connection_context: