Modify pubsub/sub_server classes for other executors

Author: hballard

graphql_subscriptions/executors/asyncio.py

@@ -1,4 +1,3 @@
-# source: https://github.com/graphql-python/graphql-core/blob/master/graphql/execution/executors/asyncio.py
 from __future__ import absolute_import
 
 import asyncio
@@ -7,6 +6,7 @@
     from asyncio import ensure_future
 except ImportError:
     # ensure_future is only implemented in Python 3.4.4+
+    # Reference: https://github.com/graphql-python/graphql-core/blob/master/graphql/execution/executors/asyncio.py
     def ensure_future(coro_or_future, loop=None):
         """Wrap a coroutine or an awaitable in a future.
         If the argument is a Future, it is returned directly.
@@ -23,8 +23,8 @@ def ensure_future(coro_or_future, loop=None):
                 del task._source_traceback[-1]
             return task
         else:
-            raise TypeError('A Future, a coroutine or an awaitable is\
-                            required')
+            raise TypeError(
+                'A Future, a coroutine or an awaitable is required')
 
 
 class AsyncioMixin(object):
@@ -44,21 +44,31 @@ def __init__(self, loop=None):
     def sleep(time):
         yield from asyncio.sleep(time)
 
+    @staticmethod
+    @asyncio.coroutine
+    def timer(callback, period):
+        while True:
+            callback()
+            yield from asyncio.sleep(period)
+
     @staticmethod
     def kill(future):
         future.cancel()
 
-    @staticmethod
-    def join(future):
+    def join(self, future):
         self.loop.run_until_complete(asyncio.wait_for(future))
 
     def join_all(self):
-        futures = self.futures
-        self.futures = []
-        self.loop.run_until_complete(asyncio.wait(futures))
+        while self.futures:
+            futures = self.futures
+            self.futures = []
+            self.loop.run_until_complete(asyncio.wait(futures))
         return futures
 
     def execute(self, fn, *args, **kwargs):
-        future = ensure_future(result, loop=self.loop)
-        self.futures.append(future)
-        return future
+        result = fn(*args, **kwargs)
+        if isinstance(result, asyncio.Future) or asyncio.iscoroutine(result):
+            future = ensure_future(result, loop=self.loop)
+            self.futures.append(future)
+            return future
+        return result

graphql_subscriptions/executors/gevent.py

@@ -19,6 +19,12 @@ def __init__(self):
     def sleep(time):
         gevent.sleep(time)
 
+    @staticmethod
+    def timer(callback, period):
+        while True:
+            callback()
+            gevent.sleep(period)
+
     @staticmethod
     def kill(greenlet):
         greenlet.kill()

graphql_subscriptions/subscription_manager/manager.py

@@ -23,7 +23,7 @@ def __init__(self, schema, pubsub, setup_funcs={}):
         self.pubsub = pubsub
         self.setup_funcs = setup_funcs
         self.subscriptions = {}
-        self.max_subscription_id = 0
+        self.max_subscription_id = 1
 
     def publish(self, trigger_name, payload):
         self.pubsub.publish(trigger_name, payload)

graphql_subscriptions/subscription_manager/pubsub.py

@@ -7,6 +7,7 @@
 import redis
 
 from ..executors.gevent import GeventExecutor
+from ..executors.asyncio import AsyncioExecutor
 
 
 class RedisPubsub(object):
@@ -17,44 +18,57 @@ def __init__(self,
                  *args,
                  **kwargs):
 
-        if hasattr(executor, 'socket'):
-            redis.connection.socket = executor.socket
-        self.coro = None
+        if executor == AsyncioExecutor:
+            try:
+                import aredis
+            except ImportError:
+                print('You need the redis_client "aredis" for use w/ asyncio')
+            redis_client = aredis
+        else:
+            redis_client = redis
+
+        if executor == GeventExecutor:
+            redis_client.connection.socket = executor.socket
+
         self.executor = executor()
+        self.get_message_task = None
 
-        self.redis = redis.StrictRedis(host, port, *args, **kwargs)
-        self.pubsub = self.redis.pubsub()
         self.subscriptions = {}
-        self.sub_id_counter = 0
+        self.sub_id_counter = 1
+
+        self.redis = redis_client.StrictRedis(host, port, *args, **kwargs)
+        self.pubsub = self.redis.pubsub()
 
     def publish(self, trigger_name, message):
-        self.redis.publish(trigger_name, pickle.dumps(message))
+        self.executor.execute(
+            self.redis.publish, trigger_name, pickle.dumps(message))
         return True
 
     def subscribe(self, trigger_name, on_message_handler, options):
         self.sub_id_counter += 1
         try:
             if trigger_name not in list(self.subscriptions.values())[0]:
-                self.pubsub.subscribe(trigger_name)
+                self.executor.execute(self.pubsub.subscribe, trigger_name)
         except IndexError:
-            self.pubsub.subscribe(trigger_name)
+            self.executor.execute(self.pubsub.subscribe, trigger_name)
         self.subscriptions[self.sub_id_counter] = [
             trigger_name, on_message_handler
         ]
-        if not self.coro:
-            self.coro = self.executor.execute(self.wait_and_get_message)
+        if not self.get_message_task:
+            self.get_message_task = self.executor.execute(
+                self.wait_and_get_message)
         return Promise.resolve(self.sub_id_counter)
 
     def unsubscribe(self, sub_id):
         trigger_name, on_message_handler = self.subscriptions[sub_id]
         del self.subscriptions[sub_id]
         try:
             if trigger_name not in list(self.subscriptions.values())[0]:
-                self.pubsub.unsubscribe(trigger_name)
+                self.executor.execute(self.pubsub.unsubscribe, trigger_name)
         except IndexError:
-            self.pubsub.unsubscribe(trigger_name)
+            self.executor.execute(self.pubsub.unsubscribe, trigger_name)
         if not self.subscriptions:
-            self.coro = self.executor.kill(self.coro)
+            self.get_message_task = self.executor.kill(self.get_message_task)
 
     def wait_and_get_message(self):
         while True:

graphql_subscriptions/subscription_transport_ws/base.py

@@ -2,7 +2,6 @@
 from promise import Promise
 import json
 
-from ..executors.gevent import GeventExecutor
 from .protocols import (SUBSCRIPTION_FAIL, SUBSCRIPTION_END, SUBSCRIPTION_DATA,
                         SUBSCRIPTION_START, SUBSCRIPTION_SUCCESS, KEEPALIVE,
                         INIT, INIT_SUCCESS, INIT_FAIL, GRAPHQL_SUBSCRIPTIONS)
@@ -38,11 +37,6 @@ def __init__(self,
 
         super(BaseSubscriptionServer, self).__init__(websocket)
 
-    def timer(self, callback, period):
-        while True:
-            callback()
-            self.executor.sleep(period)
-
     def unsubscribe(self, graphql_sub_id):
         self.subscription_manager.unsubscribe(graphql_sub_id)
 
@@ -62,7 +56,7 @@ def keep_alive_callback():
 
         if self.keep_alive:
             keep_alive_timer = self.executor.execute(
-                self.timer,
+                self.executor.timer,
                 keep_alive_callback,
                 self.keep_alive)
 
@@ -252,21 +246,22 @@ def subscription_end_promise_handler(result):
 
     def send_subscription_data(self, sub_id, payload):
         message = {'type': SUBSCRIPTION_DATA, 'id': sub_id, 'payload': payload}
-        self.ws.send(json.dumps(message))
+        self.executor.execute(self.ws.send, json.dumps(message))
 
     def send_subscription_fail(self, sub_id, payload):
         message = {'type': SUBSCRIPTION_FAIL, 'id': sub_id, 'payload': payload}
-        self.ws.send(json.dumps(message))
+        self.executor.execute(self.ws.send, json.dumps(message))
+        # self.ws.send(json.dumps(message))
 
     def send_subscription_success(self, sub_id):
         message = {'type': SUBSCRIPTION_SUCCESS, 'id': sub_id}
-        self.ws.send(json.dumps(message))
+        self.executor.execute(self.ws.send, json.dumps(message))
 
     def send_init_result(self, result):
-        self.ws.send(json.dumps(result))
+        self.executor.execute(self.ws.send, json.dumps(result))
         if result.get('type') == INIT_FAIL:
             self.ws.close(1011)
 
     def send_keep_alive(self):
         message = {'type': KEEPALIVE}
-        self.ws.send(json.dumps(message))
+        self.executor.execute(self.ws.send, json.dumps(message))

tests/test_subscription_manager.py

@@ -31,13 +31,13 @@ def test_pubsub_subscribe_and_publish(pubsub, executor, test_input, expected):
     def message_callback(message):
         try:
             assert message == expected
-            executor.kill(pubsub.coro)
+            executor.kill(pubsub.get_message_task)
         except AssertionError as e:
             sys.exit(e)
 
     def publish_callback(sub_id):
         assert pubsub.publish('a', test_input)
-        executor.join(pubsub.coro)
+        executor.join(pubsub.get_message_task)
 
     p1 = pubsub.subscribe('a', message_callback, {})
     p2 = p1.then(publish_callback)
@@ -52,7 +52,7 @@ def unsubscribe_publish_callback(sub_id):
         pubsub.unsubscribe(sub_id)
         assert pubsub.publish('a', 'test')
         try:
-            executor.join(pubsub.coro)
+            executor.join(pubsub.get_message_task)
         except AttributeError:
             return
 
@@ -190,13 +190,13 @@ def test_subscribe_with_valid_query_and_return_root_value(sub_mgr, executor):
     def callback(e, payload):
         try:
             assert payload.data.get('testSubscription') == 'good'
-            executor.kill(sub_mgr.pubsub.coro)
+            executor.kill(sub_mgr.pubsub.get_message_task)
         except AssertionError as e:
             sys.exit(e)
 
     def publish_and_unsubscribe_handler(sub_id):
         sub_mgr.publish('testSubscription', 'good')
-        executor.join(sub_mgr.pubsub.coro)
+        executor.join(sub_mgr.pubsub.get_message_task)
         sub_mgr.unsubscribe(sub_id)
 
     p1 = sub_mgr.subscribe(query, 'X', callback, {}, {}, None, None)
@@ -217,14 +217,14 @@ def callback(err, payload):
                     assert True
                 else:
                     assert payload.data.get('testFilter') == 'good_filter'
-                    executor.kill(sub_mgr.pubsub.coro)
+                    executor.kill(sub_mgr.pubsub.get_message_task)
             except AssertionError as e:
                 sys.exit(e)
 
     def publish_and_unsubscribe_handler(sub_id):
         sub_mgr.publish('filter_1', {'filterBoolean': False})
         sub_mgr.publish('filter_1', {'filterBoolean': True})
-        executor.join(sub_mgr.pubsub.coro)
+        executor.join(sub_mgr.pubsub.get_message_task)
         sub_mgr.unsubscribe(sub_id)
 
     p1 = sub_mgr.subscribe(query, 'Filter1', callback, {'filterBoolean': True},
@@ -246,15 +246,15 @@ def callback(err, payload):
                     assert True
                 else:
                     assert payload.data.get('testFilter') == 'good_filter'
-                    executor.kill(sub_mgr.pubsub.coro)
+                    executor.kill(sub_mgr.pubsub.get_message_task)
             except AssertionError as e:
                 sys.exit(e)
 
     def publish_and_unsubscribe_handler(sub_id):
         sub_mgr.publish('filter_2', {'filterBoolean': False})
         sub_mgr.publish('filter_2', {'filterBoolean': True})
         try:
-            executor.join(sub_mgr.pubsub.coro)
+            executor.join(sub_mgr.pubsub.get_message_task)
         except:
             raise
         sub_mgr.unsubscribe(sub_id)
@@ -285,13 +285,13 @@ def callback(err, payload):
             except AssertionError as e:
                 sys.exit(e)
         if non_local['trigger_count'] == 2:
-            executor.kill(sub_mgr.pubsub.coro)
+            executor.kill(sub_mgr.pubsub.get_message_task)
 
     def publish_and_unsubscribe_handler(sub_id):
         sub_mgr.publish('not_a_trigger', {'filterBoolean': False})
         sub_mgr.publish('trigger_1', {'filterBoolean': True})
         sub_mgr.publish('trigger_2', {'filterBoolean': True})
-        executor.join(sub_mgr.pubsub.coro)
+        executor.join(sub_mgr.pubsub.get_message_task)
         sub_mgr.unsubscribe(sub_id)
 
     p1 = sub_mgr.subscribe(query, 'multiTrigger', callback,
@@ -342,7 +342,7 @@ def unsubscribe_and_publish_handler(sub_id):
         sub_mgr.unsubscribe(sub_id)
         sub_mgr.publish('testSubscription', 'good')
         try:
-            executor.join(sub_mgr.pubsub.coro)
+            executor.join(sub_mgr.pubsub.get_message_task)
         except AttributeError:
             return
 
@@ -395,14 +395,14 @@ def callback(err, payload):
             assert err.message == 'Variable "$uga" of required type\
  "Boolean!" was not provided.'
 
-            executor.kill(sub_mgr.pubsub.coro)
+            executor.kill(sub_mgr.pubsub.get_message_task)
         except AssertionError as e:
             sys.exit(e)
 
     def unsubscribe_and_publish_handler(sub_id):
         sub_mgr.publish('testSubscription', 'good')
         try:
-            executor.join(sub_mgr.pubsub.coro)
+            executor.join(sub_mgr.pubsub.get_message_task)
         except AttributeError:
             return
         sub_mgr.unsubscribe(sub_id)
@@ -426,14 +426,14 @@ def callback(err, payload):
         try:
             assert err is None
             assert payload.data.get('testContext') == 'trigger'
-            executor.kill(sub_mgr.pubsub.coro)
+            executor.kill(sub_mgr.pubsub.get_message_task)
         except AssertionError as e:
             sys.exit(e)
 
     def unsubscribe_and_publish_handler(sub_id):
         sub_mgr.publish('context_trigger', 'ignored')
         try:
-            executor.join(sub_mgr.pubsub.coro)
+            executor.join(sub_mgr.pubsub.get_message_task)
         except AttributeError:
             return
         sub_mgr.unsubscribe(sub_id)
@@ -458,14 +458,14 @@ def callback(err, payload):
         try:
             assert payload is None
             assert str(err) == 'context error'
-            executor.kill(sub_mgr.pubsub.coro)
+            executor.kill(sub_mgr.pubsub.get_message_task)
         except AssertionError as e:
             sys.exit(e)
 
     def unsubscribe_and_publish_handler(sub_id):
         sub_mgr.publish('context_trigger', 'ignored')
         try:
-            executor.join(sub_mgr.pubsub.coro)
+            executor.join(sub_mgr.pubsub.get_message_task)
         except AttributeError:
             return
         sub_mgr.unsubscribe(sub_id)