Support async generator responses in django channels

Author: SmileyChris

examples/django_channels2/django_channels2/schema.py

@@ -1,5 +1,9 @@
 import graphene
 from rx import Observable
+from channels.layers import get_channel_layer
+from asgiref.sync import async_to_sync
+
+channel_layer = get_channel_layer()
 
 
 class Query(graphene.ObjectType):
@@ -9,15 +13,41 @@ def resolve_hello(self, info, **kwargs):
         return "world"
 
 
+class TestMessageMutation(graphene.Mutation):
+    class Arguments:
+        input_text = graphene.String()
+
+    output_text = graphene.String()
+
+    def mutate(self, info, input_text):
+        async_to_sync(channel_layer.group_send)("new_message", {"data": input_text})
+        return TestMessageMutation(output_text=input_text)
+
+
+class Mutations(graphene.ObjectType):
+    test_message = TestMessageMutation.Field()
+
+
 class Subscription(graphene.ObjectType):
     count_seconds = graphene.Int(up_to=graphene.Int())
+    new_message = graphene.String()
 
-    def resolve_count_seconds(root, info, up_to=5):
+    def resolve_count_seconds(self, info, up_to=5):
         return (
             Observable.interval(1000)
             .map(lambda i: "{0}".format(i))
             .take_while(lambda i: int(i) <= up_to)
         )
 
+    async def resolve_new_message(self, info):
+        channel_name = await channel_layer.new_channel()
+        await channel_layer.group_add("new_message", channel_name)
+        try:
+            while True:
+                message = await channel_layer.receive(channel_name)
+                yield message["data"]
+        finally:
+            await channel_layer.group_discard("new_message", channel_name)
+
 
-schema = graphene.Schema(query=Query, subscription=Subscription)
+schema = graphene.Schema(query=Query, mutation=Mutations, subscription=Subscription)

examples/django_channels2/django_channels2/settings.py

@@ -26,4 +26,5 @@
 ASGI_APPLICATION = "graphql_ws.django.routing.application"
 
 
+CHANNEL_LAYERS = {"default": {"BACKEND": "channels.layers.InMemoryChannelLayer"}}
 GRAPHENE = {"MIDDLEWARE": [], "SCHEMA": "django_channels2.schema.schema"}

graphql_ws/django/subscriptions.py

@@ -1,9 +1,12 @@
-from asgiref.sync import async_to_sync
+from inspect import isawaitable
 from graphene_django.settings import graphene_settings
 from graphql.execution.executors.asyncio import AsyncioExecutor
-from rx import Observer, Observable
+from rx import Observer
 from ..base import BaseConnectionContext, BaseSubscriptionServer
-from ..constants import GQL_CONNECTION_ACK, GQL_CONNECTION_ERROR
+from ..constants import GQL_CONNECTION_ACK, GQL_CONNECTION_ERROR, GQL_COMPLETE
+from ..observable_aiter import setup_observable_extension
+
+setup_observable_extension()
 
 
 class SubscriptionObserver(Observer):
@@ -76,24 +79,25 @@ async def on_connection_init(self, connection_context, op_id, payload):
             await connection_context.close(1011)
 
     async def on_start(self, connection_context, op_id, params):
-        try:
-            execution_result = await self.execute(
-                connection_context.request_context, params
+        execution_result = self.execute(connection_context.request_context, params)
+
+        if isawaitable(execution_result):
+            execution_result = await execution_result
+
+        if not hasattr(execution_result, "__aiter__"):
+            await self.send_execution_result(
+                connection_context, op_id, execution_result
             )
-            assert isinstance(
-                execution_result, Observable
-            ), "A subscription must return an observable"
-            execution_result.subscribe(
-                SubscriptionObserver(
-                    connection_context,
-                    op_id,
-                    async_to_sync(self.send_execution_result),
-                    async_to_sync(self.send_error),
-                    async_to_sync(self.on_close),
+        else:
+            iterator = await execution_result.__aiter__()
+            connection_context.register_operation(op_id, iterator)
+            async for single_result in iterator:
+                if not connection_context.has_operation(op_id):
+                    break
+                await self.send_execution_result(
+                    connection_context, op_id, single_result
                 )
-            )
-        except Exception as e:
-            self.send_error(connection_context, op_id, str(e))
+            await self.send_message(connection_context, op_id, GQL_COMPLETE)
 
     async def on_close(self, connection_context):
         remove_operations = list(connection_context.operations.keys())