Merge pull request #148 from minrk/poke-proxy-api

add proxy API

jupyterhub/apihandlers/__init__.py

@@ -1,9 +1,10 @@
 from .base import *
 from .auth import *
+from .proxy import *
 from .users import *
 
-from . import auth, users
+from . import auth, proxy, users
 
 default_handlers = []
-for mod in (auth, users):
+for mod in (auth, proxy, users):
     default_handlers.extend(mod.default_handlers)

jupyterhub/apihandlers/proxy.py

@@ -0,0 +1,68 @@
+"""Proxy handlers"""
+
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+import json
+
+from tornado import gen, web
+
+from .. import orm
+from ..utils import admin_only
+from .base import APIHandler
+
+class ProxyAPIHandler(APIHandler):
+    
+    @admin_only
+    @gen.coroutine
+    def get(self):
+        """GET /api/proxy fetches the routing table
+        
+        This is the same as fetching the routing table directly from the proxy,
+        but without clients needing to maintain separate
+        """
+        routes = yield self.proxy.get_routes()
+        self.write(json.dumps(routes))
+    
+    @admin_only
+    @gen.coroutine
+    def post(self):
+        """POST checks the proxy to ensure"""
+        yield self.proxy.check_routes()
+        
+    
+    @admin_only
+    @gen.coroutine
+    def patch(self):
+        """PATCH updates the location of the proxy
+        
+        Can be used to notify the Hub that a new proxy is in charge
+        """
+        if not self.request.body:
+            raise web.HTTPError(400, "need JSON body")
+        
+        try:
+            model = json.loads(self.request.body.decode('utf8', 'replace'))
+        except ValueError:
+            raise web.HTTPError(400, "Request body must be JSON dict")
+        if not isinstance(model, dict):
+            raise web.HTTPError(400, "Request body must be JSON dict")
+        
+        server = self.proxy.api_server
+        if 'ip' in model:
+            server.ip = model['ip']
+        if 'port' in model:
+            server.port = model['port']
+        if 'protocol' in model:
+            server.proto = model['protocol']
+        if 'auth_token' in model:
+            self.proxy.auth_token = model['auth_token']
+        self.db.commit()
+        self.log.info("Updated proxy at %s", server.url)
+        yield self.proxy.check_routes()
+        
+
+
+default_handlers = [
+    (r"/api/proxy", ProxyAPIHandler),
+]

jupyterhub/apihandlers/users.py

@@ -1,4 +1,4 @@
-"""Authorization handlers"""
+"""User handlers"""
 
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.

jupyterhub/app.py

@@ -350,6 +350,9 @@ class JupyterHub(Application):
     handlers = List()
     
     _log_formatter_cls = LogFormatter
+    http_server = None
+    proxy_process = None
+    io_loop = None
     
     def _log_level_default(self):
         return logging.INFO
@@ -869,11 +872,12 @@ class JupyterHub(Application):
             user.last_activity = max(user.last_activity, dt)
 
         self.db.commit()
+        yield self.proxy.check_routes(routes)
     
     @gen.coroutine
     def start(self):
         """Start the whole thing"""
-        loop = IOLoop.current()
+        self.io_loop = loop = IOLoop.current()
         
         if self.subapp:
             self.subapp.start()
@@ -907,12 +911,19 @@ class JupyterHub(Application):
             pc.start()
 
         # start the webserver
-        http_server = tornado.httpserver.HTTPServer(self.tornado_application, xheaders=True)
-        http_server.listen(self.hub_port)
+        self.http_server = tornado.httpserver.HTTPServer(self.tornado_application, xheaders=True)
+        self.http_server.listen(self.hub_port)
         # run the cleanup step (in a new loop, because the interrupted one is unclean)
         
         atexit.register(lambda : IOLoop().run_sync(self.cleanup))
     
+    def stop(self):
+        if not self.io_loop:
+            return
+        if self.http_server:
+            self.io_loop.add_callback(self.http_server.stop)
+        self.io_loop.add_callback(self.io_loop.stop)
+    
     @gen.coroutine
     def launch_instance_async(self, argv=None):
         yield self.initialize(argv)

jupyterhub/orm.py

@@ -193,6 +193,23 @@ class Proxy(Base):
         resp = yield self.api_request('', client=client)
         return json.loads(resp.body.decode('utf8', 'replace'))
 
+    @gen.coroutine
+    def check_routes(self, routes=None):
+        """Check that all users are properly"""
+        if not routes:
+            routes = yield self.get_routes()
+
+        have_routes = { r['user'] for r in routes if 'user' in r }
+        futures = []
+        db = inspect(self).session
+        for user in db.query(User).filter(User.server != None):
+            if user.name not in have_routes:
+                self.log.warn("Adding missing route for %s", user.name)
+                futures.append(self.add_user(user))
+        for f in futures:
+            yield f
+
+
 
 class Hub(Base):
     """Bring it all together at the hub.

jupyterhub/tests/conftest.py

@@ -47,6 +47,10 @@ def io_loop():
 @fixture(scope='module')
 def app(request):
     app = MockHub.instance(log_level=logging.DEBUG)
+    print(app)
     app.start([])
-    request.addfinalizer(app.stop)
+    def fin():
+        MockHub.clear_instance()
+        app.stop()
+    request.addfinalizer(fin)
     return app

jupyterhub/tests/mocking.py

@@ -120,7 +120,7 @@ class MockHub(JupyterHub):
         evt.wait(timeout=5)
     
     def stop(self):
-        self.io_loop.add_callback(self.io_loop.stop)
+        super().stop()
         self._thread.join()
         IOLoop().run_sync(self.cleanup)
         # ignore the call that will fire in atexit

jupyterhub/tests/test_api.py

@@ -277,3 +277,10 @@ def test_never_spawn(app, io_loop):
     assert not user.spawn_pending
     status = io_loop.run_sync(user.spawner.poll)
     assert status is not None
+
+
+def test_get_proxy(app, io_loop):
+    r = api_request(app, 'proxy')
+    r.raise_for_status()
+    reply = r.json()
+    assert list(reply.keys()) == ['/']

jupyterhub/tests/test_proxy.py

@@ -0,0 +1,104 @@
+"""Test a proxy being started before the Hub"""
+
+import json
+import os
+from subprocess import Popen
+
+from .mocking import MockHub
+from .test_api import api_request
+from ..utils import wait_for_http_server
+
+def test_external_proxy(request, io_loop):
+    """Test a proxy started before the Hub"""
+    auth_token='secret!'
+    proxy_ip = '127.0.0.1'
+    proxy_port = 54321
+    
+    app = MockHub.instance(
+        proxy_api_ip=proxy_ip,
+        proxy_api_port=proxy_port,
+        proxy_auth_token=auth_token,
+    )
+    request.addfinalizer(app.stop)
+    env = os.environ.copy()
+    env['CONFIGPROXY_AUTH_TOKEN'] = auth_token
+    cmd = [app.proxy_cmd,
+        '--ip', app.ip,
+        '--port', str(app.port),
+        '--api-ip', proxy_ip,
+        '--api-port', str(proxy_port),
+        '--default-target', 'http://%s:%i' % (app.hub_ip, app.hub_port),
+    ]
+    proxy = Popen(cmd, env=env)
+    def _cleanup_proxy():
+        if proxy.poll() is None:
+            proxy.terminate()
+    request.addfinalizer(_cleanup_proxy)
+    
+    def wait_for_proxy():
+        io_loop.run_sync(lambda : wait_for_http_server(
+            'http://%s:%i' % (proxy_ip, proxy_port)))
+    wait_for_proxy()
+    
+    app.start([])
+    
+    assert app.proxy_process is None
+    
+    routes = io_loop.run_sync(app.proxy.get_routes)
+    assert list(routes.keys()) == ['/']
+    
+    # add user
+    name = 'river'
+    r = api_request(app, 'users', name, method='post')
+    r.raise_for_status()
+    r = api_request(app, 'users', name, 'server', method='post')
+    r.raise_for_status()
+    
+    routes = io_loop.run_sync(app.proxy.get_routes)
+    assert sorted(routes.keys()) == ['/', '/user/river']
+    
+    # teardown the proxy and start a new one in the same place
+    proxy.terminate()
+    proxy = Popen(cmd, env=env)
+    wait_for_proxy()
+    
+    routes = io_loop.run_sync(app.proxy.get_routes)
+    assert list(routes.keys()) == ['/']
+    
+    # poke the server to update the proxy
+    r = api_request(app, 'proxy', method='post')
+    r.raise_for_status()
+    
+    # check that the routes are correct
+    routes = io_loop.run_sync(app.proxy.get_routes)
+    assert sorted(routes.keys()) == ['/', '/user/river']
+    
+    # teardown the proxy again, and start a new one with different auth and port
+    proxy.terminate()
+    new_auth_token = 'different!'
+    env['CONFIGPROXY_AUTH_TOKEN'] = new_auth_token
+    proxy_port = 55432
+    cmd = [app.proxy_cmd,
+        '--ip', app.ip,
+        '--port', str(app.port),
+        '--api-ip', app.proxy_api_ip,
+        '--api-port', str(proxy_port),
+        '--default-target', 'http://%s:%i' % (app.hub_ip, app.hub_port),
+    ]
+    
+    proxy = Popen(cmd, env=env)
+    wait_for_proxy()
+    
+    # tell the hub where the new proxy is
+    r = api_request(app, 'proxy', method='patch', data=json.dumps({
+        'port': proxy_port,
+        'auth_token': new_auth_token,
+    }))
+    r.raise_for_status()
+    assert app.proxy.api_server.port == proxy_port
+    assert app.proxy.auth_token == new_auth_token
+    
+    # check that the routes are correct
+    routes = io_loop.run_sync(app.proxy.get_routes)
+    assert sorted(routes.keys()) == ['/', '/user/river']
+