move some utility functions to ORM classes

simplify handler methods

jupyterhub/app.py

@@ -364,6 +364,7 @@ class JupyterHubApp(Application):
             config=self.config,
             log=self.log,
             db=self.db,
+            proxy=self.proxy,
             hub=self.hub,
             admin_users=self.admin_users,
             authenticator=import_item(self.authenticator)(config=self.config),

jupyterhub/handlers/base.py

@@ -3,7 +3,6 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
-import json
 import re
 
 try:
@@ -12,8 +11,6 @@ try:
 except ImportError:
     from httplib import responses
 
-import requests
-
 from jinja2 import TemplateNotFound
 
 from tornado.log import app_log
@@ -23,11 +20,12 @@ from tornado import gen, web
 
 from .. import orm
 from ..spawner import LocalProcessSpawner
-from ..utils import wait_for_server, url_path_join
+from ..utils import url_path_join
 
 # pattern for the authentication token header
 auth_header_pat = re.compile(r'^token\s+([^\s]+)$')
 
+
 class BaseHandler(RequestHandler):
     """Base Handler class with access to common methods and properties."""
 
@@ -52,6 +50,10 @@ class BaseHandler(RequestHandler):
     def hub(self):
         return self.settings['hub']
     
+    @property
+    def proxy(self):
+        return self.settings['proxy']
+    
     @property
     def authenticator(self):
         return self.settings.get('authenticator', None)
@@ -159,78 +161,20 @@ class BaseHandler(RequestHandler):
     def spawner_class(self):
         return self.settings.get('spawner_class', LocalProcessSpawner)
 
-    @gen.coroutine
-    def notify_proxy(self, user):
-        proxy = self.db.query(orm.Proxy).first()
-        r = requests.post(
-            url_path_join(
-                proxy.api_server.url,
-                user.server.base_url,
-            ),
-            data=json.dumps(dict(
-                target=user.server.host,
-                user=user.name,
-            )),
-            headers={'Authorization': "token %s" % proxy.auth_token},
-        )
-        yield wait_for_server(user.server.ip, user.server.port)
-        r.raise_for_status()
-
-    @gen.coroutine
-    def notify_proxy_delete(self, user):
-        proxy = self.db.query(orm.Proxy).first()
-        r = requests.delete(
-            url_path_join(
-                proxy.api_server.url,
-                user.server.base_url,
-            ),
-            headers={'Authorization': "token %s" % proxy.auth_token},
-        )
-        r.raise_for_status()
-
     @gen.coroutine
     def spawn_single_user(self, user):
-        user.server = orm.Server(
-            cookie_name='%s-%s' % (self.hub.server.cookie_name, user.name),
-            cookie_secret=self.hub.server.cookie_secret,
-            base_url=url_path_join(self.base_url, 'user', user.name),
-        )
-        self.db.add(user.server)
-        self.db.commit()
-
-        api_token = user.new_api_token()
-        self.db.add(api_token)
-        self.db.commit()
-
-        spawner = user.spawner = self.spawner_class(
-            config=self.config,
-            user=user,
+        yield user.spawn(
+            spawner_class=self.spawner_class,
+            base_url=self.base_url,
             hub=self.hub,
-            api_token=api_token.token,
         )
-        yield spawner.start()
-
-        # store state
-        user.state = spawner.get_state()
-        self.db.commit()
-
-        yield self.notify_proxy(user)
+        yield self.proxy.add_user(user)
         raise gen.Return(user)
     
     @gen.coroutine
     def stop_single_user(self, user):
-        if user.spawner is None:
-            return
-        status = yield user.spawner.poll()
-        if status is None:
-            yield user.spawner.stop()
-        self.notify_proxy_delete(user)
-        user.state = {}
-        user.spawner = None
-        user.server = None
-        self.db.commit()
-        
-        raise gen.Return(user)
+        yield self.proxy.delete_user(user)
+        yield user.stop()
 
     #---------------------------------------------------------------
     # template rendering

jupyterhub/orm.py

@@ -6,8 +6,12 @@
 import json
 import uuid
 
+import requests
+from tornado import gen
+
 from sqlalchemy.types import TypeDecorator, VARCHAR
 from sqlalchemy import (
+    inspect,
     Column, Integer, String, ForeignKey, Unicode, Binary, Boolean,
 )
 from sqlalchemy.ext.declarative import declarative_base, declared_attr
@@ -17,7 +21,7 @@ from sqlalchemy import create_engine
 
 from IPython.utils.py3compat import str_to_unicode
 
-from .utils import random_port, url_path_join
+from .utils import random_port, url_path_join, wait_for_server
 
 
 def new_token(*args, **kwargs):
@@ -85,6 +89,11 @@ class Server(Base):
             host=self.host,
             uri=self.base_url,
         )
+    
+    @gen.coroutine
+    def wait_up(self, timeout=10):
+        """Wait for this server to come up"""
+        yield wait_for_server(self.ip or 'localhost', self.port, timeout=timeout)
 
 
 class Proxy(Base):
@@ -109,6 +118,34 @@ class Proxy(Base):
         else:
             return "<%s [unconfigured]>" % self.__class__.__name__
 
+    @gen.coroutine
+    def add_user(self, user):
+        """Add a user's server to the proxy table."""
+        r = requests.post(
+            url_path_join(
+                self.api_server.url,
+                user.server.base_url,
+            ),
+            data=json.dumps(dict(
+                target=user.server.host,
+                user=user.name,
+            )),
+            headers={'Authorization': "token %s" % self.auth_token},
+        )
+        r.raise_for_status()
+
+    @gen.coroutine
+    def delete_user(self, user):
+        """Remove a user's server to the proxy table."""
+        r = requests.delete(
+            url_path_join(
+                self.api_server.url,
+                user.server.base_url,
+            ),
+            headers={'Authorization': "token %s" % self.auth_token},
+        )
+        r.raise_for_status()
+
 
 class Hub(Base):
     """Bring it all together at the hub.
@@ -190,6 +227,50 @@ class User(Base):
         """Return a new cookie token"""
         return self._new_token(CookieToken)
 
+    @gen.coroutine
+    def spawn(self, spawner_class, base_url='/', hub=None, config=None):
+        db = inspect(self).session
+        if hub is None:
+            hub = db.query(Hub).first()
+        self.server = Server(
+            cookie_name='%s-%s' % (hub.server.cookie_name, self.name),
+            cookie_secret=hub.server.cookie_secret,
+            base_url=url_path_join(base_url, 'user', self.name),
+        )
+        db.add(self.server)
+        db.commit()
+
+        api_token = self.new_api_token()
+        db.add(api_token)
+        db.commit()
+
+        spawner = self.spawner = spawner_class(
+            config=config,
+            user=self,
+            hub=hub,
+            api_token=api_token.token,
+        )
+        yield spawner.start()
+
+        # store state
+        self.state = spawner.get_state()
+        db.commit()
+    
+        yield self.server.wait_up()
+        raise gen.Return(self)
+
+    @gen.coroutine
+    def stop(self):
+        if self.spawner is None:
+            return
+        status = yield self.spawner.poll()
+        if status is None:
+            yield self.spawner.stop()
+        self.state = {}
+        self.spawner = None
+        self.server = None
+        inspect(self).session.commit()
+
 
 class Token(object):
     """Mixin for token tables, since we have two"""