Split base classes into sync and async

Deduplicate code

Author: SmileyChris

graphql_ws/__init__.py

@@ -5,6 +5,3 @@
 __author__ = """Syrus Akbary"""
 __email__ = 'me@syrusakbary.com'
 __version__ = '0.3.1'
-
-
-from .base import BaseConnectionContext, BaseSubscriptionServer  # noqa: F401

graphql_ws/aiohttp.py

@@ -1,23 +1,13 @@
-from inspect import isawaitable
-from asyncio import ensure_future, wait, shield
+import json
+from asyncio import ensure_future, shield
 
 from aiohttp import WSMsgType
-from graphql.execution.executors.asyncio import AsyncioExecutor
 
-from .base import (
-    ConnectionClosedException, BaseConnectionContext, BaseSubscriptionServer)
-from .observable_aiter import setup_observable_extension
+from .base import ConnectionClosedException
+from .base_async import BaseAsyncConnectionContext, BaseAsyncSubscriptionServer
 
-from .constants import (
-    GQL_CONNECTION_ACK,
-    GQL_CONNECTION_ERROR,
-    GQL_COMPLETE
-)
 
-setup_observable_extension()
-
-
-class AiohttpConnectionContext(BaseConnectionContext):
+class AiohttpConnectionContext(BaseAsyncConnectionContext):
     async def receive(self):
         msg = await self.ws.receive()
         if msg.type == WSMsgType.TEXT:
@@ -32,7 +22,7 @@ async def receive(self):
     async def send(self, data):
         if self.closed:
             return
-        await self.ws.send_str(data)
+        await self.ws.send_str(json.dumps(data))
 
     @property
     def closed(self):
@@ -42,81 +32,24 @@ async def close(self, code):
         await self.ws.close(code=code)
 
 
-class AiohttpSubscriptionServer(BaseSubscriptionServer):
-    def __init__(self, schema, keep_alive=True, loop=None):
-        self.loop = loop
-        super().__init__(schema, keep_alive)
-
-    def get_graphql_params(self, *args, **kwargs):
-        params = super(AiohttpSubscriptionServer,
-                       self).get_graphql_params(*args, **kwargs)
-        return dict(params, return_promise=True,
-                    executor=AsyncioExecutor(loop=self.loop))
-
+class AiohttpSubscriptionServer(BaseAsyncSubscriptionServer):
     async def _handle(self, ws, request_context=None):
         connection_context = AiohttpConnectionContext(ws, request_context)
         await self.on_open(connection_context)
-        pending = set()
         while True:
             try:
                 if connection_context.closed:
                     raise ConnectionClosedException()
                 message = await connection_context.receive()
             except ConnectionClosedException:
                 break
-            finally:
-                if pending:
-                    (_, pending) = await wait(pending, timeout=0, loop=self.loop)
 
-            task = ensure_future(
-                self.on_message(connection_context, message), loop=self.loop)
-            pending.add(task)
-
-        self.on_close(connection_context)
-        for task in pending:
-            task.cancel()
+            connection_context.remember_task(
+                ensure_future(
+                    self.on_message(connection_context, message), loop=self.loop
+                )
+            )
+        await self.on_close(connection_context)
 
     async def handle(self, ws, request_context=None):
         await shield(self._handle(ws, request_context), loop=self.loop)
-
-    async def on_open(self, connection_context):
-        pass
-
-    def on_close(self, connection_context):
-        remove_operations = list(connection_context.operations.keys())
-        for op_id in remove_operations:
-            self.unsubscribe(connection_context, op_id)
-
-    async def on_connect(self, connection_context, payload):
-        pass
-
-    async def on_connection_init(self, connection_context, op_id, payload):
-        try:
-            await self.on_connect(connection_context, payload)
-            await self.send_message(connection_context, op_type=GQL_CONNECTION_ACK)
-        except Exception as e:
-            await self.send_error(connection_context, op_id, e, GQL_CONNECTION_ERROR)
-            await connection_context.close(1011)
-
-    async def on_start(self, connection_context, op_id, 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)
-        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)
-            await self.send_message(connection_context, op_id, GQL_COMPLETE)
-
-    async def on_stop(self, connection_context, op_id):
-        self.unsubscribe(connection_context, op_id)

graphql_ws/base.py

@@ -1,16 +1,16 @@
 import json
 from collections import OrderedDict
 
-from graphql import graphql, format_error
+from graphql import format_error
 
 from .constants import (
+    GQL_CONNECTION_ERROR,
     GQL_CONNECTION_INIT,
     GQL_CONNECTION_TERMINATE,
+    GQL_DATA,
+    GQL_ERROR,
     GQL_START,
     GQL_STOP,
-    GQL_ERROR,
-    GQL_CONNECTION_ERROR,
-    GQL_DATA
 )
 
 
@@ -51,33 +51,16 @@ def close(self, code):
 
 
 class BaseSubscriptionServer(object):
+    graphql_executor = None
 
     def __init__(self, schema, keep_alive=True):
         self.schema = schema
         self.keep_alive = keep_alive
 
-    def get_graphql_params(self, connection_context, payload):
-        return {
-            'request_string': payload.get('query'),
-            'variable_values': payload.get('variables'),
-            'operation_name': payload.get('operationName'),
-            'context_value': payload.get('context'),
-        }
-
-    def build_message(self, id, op_type, payload):
-        message = {}
-        if id is not None:
-            message['id'] = id
-        if op_type is not None:
-            message['type'] = op_type
-        if payload is not None:
-            message['payload'] = payload
-        return message
-
     def process_message(self, connection_context, parsed_message):
-        op_id = parsed_message.get('id')
-        op_type = parsed_message.get('type')
-        payload = parsed_message.get('payload')
+        op_id = parsed_message.get("id")
+        op_type = parsed_message.get("type")
+        payload = parsed_message.get("payload")
 
         if op_type == GQL_CONNECTION_INIT:
             return self.on_connection_init(connection_context, op_id, payload)
@@ -92,22 +75,63 @@ def process_message(self, connection_context, parsed_message):
             if not isinstance(params, dict):
                 error = Exception(
                     "Invalid params returned from get_graphql_params!"
-                    " Return values must be a dict.")
+                    " Return values must be a dict."
+                )
                 return self.send_error(connection_context, op_id, error)
 
             # If we already have a subscription with this id, unsubscribe from
             # it first
             if connection_context.has_operation(op_id):
                 self.unsubscribe(connection_context, op_id)
 
+            params = self.get_graphql_params(connection_context, payload)
             return self.on_start(connection_context, op_id, params)
 
         elif op_type == GQL_STOP:
             return self.on_stop(connection_context, op_id)
 
         else:
-            return self.send_error(connection_context, op_id, Exception(
-                "Invalid message type: {}.".format(op_type)))
+            return self.send_error(
+                connection_context,
+                op_id,
+                Exception("Invalid message type: {}.".format(op_type)),
+            )
+
+    def on_connection_init(self, connection_context, op_id, payload):
+        raise NotImplementedError("on_connection_init method not implemented")
+
+    def on_connection_terminate(self, connection_context, op_id):
+        return connection_context.close(1011)
+
+    def get_graphql_params(self, connection_context, payload):
+        return {
+            "request_string": payload.get("query"),
+            "variable_values": payload.get("variables"),
+            "operation_name": payload.get("operationName"),
+            "context_value": payload.get("context"),
+            "executor": self.graphql_executor(),
+        }
+
+    def on_open(self, connection_context):
+        raise NotImplementedError("on_open method not implemented")
+
+    def on_stop(self, connection_context, op_id):
+        raise NotImplementedError("on_stop method not implemented")
+
+    def send_message(self, connection_context, op_id=None, op_type=None, payload=None):
+        message = self.build_message(op_id, op_type, payload)
+        assert message, "You need to send at least one thing"
+        return connection_context.send(message)
+
+    def build_message(self, id, op_type, payload):
+        message = {}
+        if id is not None:
+            message["id"] = id
+        if op_type is not None:
+            message["type"] = op_type
+        if payload is not None:
+            message["payload"] = payload
+        return message
 
     def send_execution_result(self, connection_context, op_id, execution_result):
         result = self.execution_result_to_dict(execution_result)
@@ -116,86 +140,34 @@ def send_execution_result(self, connection_context, op_id, execution_result):
     def execution_result_to_dict(self, execution_result):
         result = OrderedDict()
         if execution_result.data:
-            result['data'] = execution_result.data
+            result["data"] = execution_result.data
         if execution_result.errors:
-            result['errors'] = [format_error(error)
-                                for error in execution_result.errors]
+            result["errors"] = [
+                format_error(error) for error in execution_result.errors
+            ]
         return result
 
-    def send_message(self, connection_context, op_id=None, op_type=None, payload=None):
-        message = self.build_message(op_id, op_type, payload)
-        assert message, "You need to send at least one thing"
-        json_message = json.dumps(message)
-        return connection_context.send(json_message)
-
     def send_error(self, connection_context, op_id, error, error_type=None):
         if error_type is None:
             error_type = GQL_ERROR
 
         assert error_type in [GQL_CONNECTION_ERROR, GQL_ERROR], (
-            'error_type should be one of the allowed error messages'
-            ' GQL_CONNECTION_ERROR or GQL_ERROR'
-        )
-
-        error_payload = {
-            'message': str(error)
-        }
-
-        return self.send_message(
-            connection_context,
-            op_id,
-            error_type,
-            error_payload
+            "error_type should be one of the allowed error messages"
+            " GQL_CONNECTION_ERROR or GQL_ERROR"
         )
 
-    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)
+        error_payload = {"message": str(error)}
 
-    def on_operation_complete(self, connection_context, op_id):
-        pass
-
-    def on_connection_terminate(self, connection_context, op_id):
-        return connection_context.close(1011)
-
-    def execute(self, request_context, params):
-        return graphql(
-            self.schema, **dict(params, allow_subscriptions=True))
-
-    def handle(self, ws, request_context=None):
-        raise NotImplementedError("handle method not implemented")
+        return self.send_message(connection_context, op_id, error_type, error_payload)
 
     def on_message(self, connection_context, message):
         try:
             if not isinstance(message, dict):
                 parsed_message = json.loads(message)
-                assert isinstance(
-                    parsed_message, dict), "Payload must be an object."
+                assert isinstance(parsed_message, dict), "Payload must be an object."
             else:
                 parsed_message = message
         except Exception as e:
             return self.send_error(connection_context, None, e)
 
         return self.process_message(connection_context, parsed_message)
-
-    def on_open(self, connection_context):
-        raise NotImplementedError("on_open method not implemented")
-
-    def on_connect(self, connection_context, payload):
-        raise NotImplementedError("on_connect method not implemented")
-
-    def on_close(self, connection_context):
-        raise NotImplementedError("on_close method not implemented")
-
-    def on_connection_init(self, connection_context, op_id, payload):
-        raise NotImplementedError("on_connection_init method not implemented")
-
-    def on_stop(self, connection_context, op_id):
-        raise NotImplementedError("on_stop method not implemented")
-
-    def on_start(self, connection_context, op_id, params):
-        raise NotImplementedError("on_start method not implemented")