Add SubscriptionHandler

cylc/uiserver/handlers.py

@@ -17,14 +17,17 @@
 import json
 import os
 import re
+from asyncio import Queue
 from subprocess import Popen, PIPE
 from typing import List, Union
 
-from jupyterhub import __version__ as jupyterhub_version
-from jupyterhub.services.auth import HubOAuthenticated
-from tornado import web
 from graphene_tornado.tornado_graphql_handler import TornadoGraphQLHandler
 from graphql import get_default_backend
+from graphql_ws.constants import GRAPHQL_WS
+from jupyterhub import __version__ as jupyterhub_version
+from jupyterhub.services.auth import HubOAuthenticated
+from tornado import web, websocket
+from tornado.ioloop import IOLoop
 
 
 class BaseHandler(web.RequestHandler):
@@ -145,9 +148,32 @@ def prepare(self):
         super().prepare()
 
 
+class SubscriptionHandler(websocket.WebSocketHandler):
+
+    def initialize(self, sub_server):
+        self.subscription_server = sub_server
+        self.queue = Queue(100)
+
+    def select_subprotocol(self, subprotocols):
+        return GRAPHQL_WS
+
+    def open(self, *args, **kwargs):
+        IOLoop.current().spawn_callback(self.subscription_server.handle, self)
+
+    async def on_message(self, message):
+        await self.queue.put(message)
+
+    async def recv(self):
+        return await self.queue.get()
+
+    def check_origin(self, origin: str) -> bool:
+        return True
+
+
 __all__ = [
     "MainHandler",
     "UserProfileHandler",
     "CylcScanHandler",
-    "UIServerGraphQLHandler"
+    "UIServerGraphQLHandler",
+    "SubscriptionHandler"
 ]

cylc/uiserver/main.py

@@ -26,6 +26,7 @@
 from logging.config import dictConfig
 from os.path import join, abspath, dirname
 from tornado import web, ioloop
+from tornado_ws import TornadoSubscriptionServer
 
 from jupyterhub.services.auth import HubOAuthCallbackHandler
 from jupyterhub.utils import url_path_join
@@ -119,6 +120,10 @@ def _make_app(self, debug: bool):
                     UIServerGraphQLHandler,
                     dict(schema=schema, resolvers=self.resolvers,
                          graphiql=True)),
+                (url_path_join(
+                    self._jupyter_hub_service_prefix, '/subscriptions'),
+                    SubscriptionHandler,
+                    dict(sub_server=TornadoSubscriptionServer(schema)))
             ],
             # FIXME: decide (and document) whether cookies will be permanent
             # after server restart.