Added request context for each server handle

Author: syrusakbary

graphql_ws/aiohttp.py

@@ -39,10 +39,10 @@ class AiohttpSubscriptionServer(BaseSubscriptionServer):
     def get_graphql_params(self, *args, **kwargs):
         params = super(AiohttpSubscriptionServer,
                        self).get_graphql_params(*args, **kwargs)
-        return dict(params, executor=AsyncioExecutor())
+        return dict(params, return_promise=True, executor=AsyncioExecutor())
 
-    async def handle(self, ws):
-        connection_context = AiohttpConnectionContext(ws)
+    async def handle(self, ws, request_context=None):
+        connection_context = AiohttpConnectionContext(ws, request_context)
         await self.on_open(connection_context)
         while True:
             try:
@@ -70,19 +70,13 @@ 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)
-
-            # if self.keep_alive:
-            # await self.send_message(connection_context,
-            # op_type=GQL_CONNECTION_KEEP_ALIVE)
         except Exception as e:
             await self.send_error(connection_context, op_id, e, GQL_CONNECTION_ERROR)
             await connection_context.close(1011)
 
-    async def on_connection_terminate(self, connection_context, op_id):
-        await connection_context.close(1011)
-
     async def on_start(self, connection_context, op_id, params):
-        execution_result = self.execute(return_promise=True, **params)
+        execution_result = self.execute(
+            connection_context.request_context, params)
 
         if isawaitable(execution_result):
             execution_result = await execution_result

graphql_ws/base.py

@@ -21,9 +21,10 @@ class ConnectionClosedException(Exception):
 
 
 class BaseConnectionContext(object):
-    def __init__(self, ws):
+    def __init__(self, ws, request_context=None):
         self.ws = ws
         self.operations = {}
+        self.request_context = request_context
 
     def has_operation(self, op_id):
         return op_id in self.operations
@@ -167,13 +168,13 @@ def on_operation_complete(self, connection_context, op_id):
         pass
 
     def on_connection_terminate(self, connection_context, op_id):
-        return connection_context.close()
+        return connection_context.close(1011)
 
-    def execute(self, **params):
+    def execute(self, request_context, params):
         return graphql(
             self.schema, **dict(params, allow_subscriptions=True))
 
-    def handle(self, ws):
+    def handle(self, ws, request_context=None):
         raise NotImplementedError("handle method not implemented")
 
     def on_message(self, connection_context, message):

graphql_ws/gevent.py

@@ -42,8 +42,8 @@ def get_graphql_params(self, *args, **kwargs):
                        self).get_graphql_params(*args, **kwargs)
         return dict(params, executor=SyncExecutor())
 
-    def handle(self, ws):
-        connection_context = GeventConnectionContext(ws)
+    def handle(self, ws, request_context=None):
+        connection_context = GeventConnectionContext(ws, request_context)
         self.on_open(connection_context)
         while True:
             try:
@@ -75,12 +75,10 @@ def on_connection_init(self, connection_context, op_id, payload):
             self.send_error(connection_context, op_id, e, GQL_CONNECTION_ERROR)
             connection_context.close(1011)
 
-    def on_connection_terminate(self, connection_context, op_id):
-        connection_context.close(1011)
-
     def on_start(self, connection_context, op_id, params):
         try:
-            execution_result = self.execute(**params)
+            execution_result = self.execute(
+                connection_context.request_context, params)
             assert isinstance(
                 execution_result, Observable), "A subscription must return an observable"
             execution_result.subscribe(SubscriptionObserver(

setup.py

@@ -36,7 +36,7 @@
     long_description=readme + '\n\n' + history,
     author="Syrus Akbary",
     author_email='me@syrusakbary.com',
-    url='https://github.com/graphql-python/graphql_ws',
+    url='https://github.com/graphql-python/graphql-ws',
     packages=find_packages(include=['graphql_ws']),
     include_package_data=True,
     install_requires=requirements,