fix graphql send cancel on break from subscription loop

Author: jell-o-fishi

rsocket/graphql/rsocket_transport.py

@@ -9,6 +9,7 @@
 from reactivestreams.subscriber import DefaultSubscriber
 from rsocket.extensions.helpers import composite, route
 from rsocket.frame_helpers import str_to_bytes
+from rsocket.logger import logger
 from rsocket.payload import Payload
 from rsocket.rsocket_client import RSocketClient
 
@@ -104,6 +105,14 @@ def on_complete(self):
             def on_error(self, exception: Exception):
                 self._received_queue.put_nowait(exception)
 
+            def cancel(self):
+                if self.subscription is not None:
+                    self.subscription.cancel()
+
+            def request(self, n: int):
+                if self.subscription is not None:
+                    self.subscription.request(n)
+
         rsocket_payload = self._create_rsocket_payload(document, variable_values, operation_name)
 
         received_queue = Queue()
@@ -120,4 +129,12 @@ def on_error(self, exception: Exception):
             if response is complete_object:
                 break
 
-            yield self._response_to_execution_result(response)
+            execution_result = self._response_to_execution_result(response)
+
+            try:
+                yield execution_result
+            except GeneratorExit:
+                logger().debug('Generator exited')
+                subscriber.cancel()
+                return
+

tests/test_integrations/test_graphql.py

@@ -1,3 +1,5 @@
+import asyncio
+
 from gql import Client, gql
 
 from rsocket.extensions.mimetypes import WellKnownMimeTypes
@@ -75,4 +77,38 @@ def handler_factory():
             gql("""{__schema { types { name  }  } }"""),
             get_execution_result=True)
 
-        assert response.data == expected_schema
+        assert response.data == expected_schema
+
+
+async def test_graphql_break_loop(lazy_pipe, graphql_schema):
+    def handler_factory():
+        return RoutingRequestHandler(graphql_handler(graphql_schema, 'graphql'))
+
+    async with lazy_pipe(
+            client_arguments={'metadata_encoding': WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA},
+            server_arguments={'handler_factory': handler_factory}) as (server, client):
+        graphql = Client(
+            schema=graphql_schema,
+            transport=RSocketTransport(client),
+        )
+
+        responses = []
+        i = 0
+        async for response in graphql.subscribe_async(
+                document=gql("""
+                subscription {
+                    greetings {message}
+                }
+                """),
+                get_execution_result=True):
+            responses.append(response.data)
+            i += 1
+            if i > 4:
+                break
+
+        assert len(responses) == 5
+        assert responses[0] == {'greetings': {'message': 'Hello world 0'}}
+
+        await asyncio.sleep(1)
+
+        # assert responses[9] == {'greetings': {'message': 'Hello world 9'}}