Move operation unsubscription to BaseSubscriptionServer

Author: SmileyChris

graphql_ws/base.py

@@ -168,3 +168,11 @@ def on_message(self, connection_context, message):
             return self.send_error(connection_context, None, e)
 
         return self.process_message(connection_context, parsed_message)
+
+    def unsubscribe(self, connection_context, op_id):
+        if connection_context.has_operation(op_id):
+            # Close async iterator
+            connection_context.get_operation(op_id).dispose()
+            # Close operation
+            connection_context.remove_operation(op_id)
+        self.on_operation_complete(connection_context, op_id)

graphql_ws/base_async.py

@@ -111,8 +111,5 @@ async def on_close(self, connection_context):
     async def on_stop(self, connection_context, op_id):
         await self.unsubscribe(connection_context, op_id)
 
-    async def unsubscribe(self, connection_context, op_id):
-        await super().unsubscribe(connection_context, op_id)
-
     async def on_operation_complete(self, connection_context, op_id):
         pass

graphql_ws/base_sync.py

@@ -8,14 +8,6 @@
 class BaseSyncSubscriptionServer(BaseSubscriptionServer):
     graphql_executor = SyncExecutor
 
-    def unsubscribe(self, connection_context, op_id):
-        if connection_context.has_operation(op_id):
-            # Close async iterator
-            connection_context.get_operation(op_id).dispose()
-            # Close operation
-            connection_context.remove_operation(op_id)
-        self.on_operation_complete(connection_context, op_id)
-
     def on_operation_complete(self, connection_context, op_id):
         pass
 
@@ -51,7 +43,7 @@ def on_start(self, connection_context, op_id, params):
             assert isinstance(
                 execution_result, Observable
             ), "A subscription must return an observable"
-            execution_result.subscribe(
+            disposable = execution_result.subscribe(
                 SubscriptionObserver(
                     connection_context,
                     op_id,
@@ -60,6 +52,8 @@ def on_start(self, connection_context, op_id, params):
                     self.on_close,
                 )
             )
+            connection_context.register_operation(op_id, disposable)
+
         except Exception as e:
             self.send_error(connection_context, op_id, str(e))