Merge pull request #48 from minrk/resume-from-db

enable resuming Hub state from database

jupyterhub/apihandlers/users.py

@@ -133,9 +133,11 @@ class UserServerAPIHandler(BaseUserHandler):
     def post(self, name):
         user = self.find_user(name)
         if user.spawner:
-            raise web.HTTPError(400, "%s's server is already running" % name)
+            state = yield user.spawner.poll()
-        else:
+            if state is None:
-            yield self.spawn_single_user(user)
+                raise web.HTTPError(400, "%s's server is already running" % name)
+
+        yield self.spawn_single_user(user)
         self.set_status(201)
 
     @gen.coroutine

jupyterhub/app.py

@@ -25,11 +25,10 @@ from tornado.log import LogFormatter
 from tornado import gen, web
 
 from IPython.utils.traitlets import (
-    Unicode, Integer, Dict, TraitError, List, Bool, Bytes, Any,
+    Unicode, Integer, Dict, TraitError, List, Bool, Any,
-    DottedObjectName, Set,
+    Type, Set, Instance,
 )
 from IPython.config import Application, catch_config_error
-from IPython.utils.importstring import import_item
 
 here = os.path.dirname(__file__)
 
@@ -37,7 +36,7 @@ from . import handlers, apihandlers
 
 from . import orm
 from ._data import DATA_FILES_PATH
-from .utils import url_path_join
+from .utils import url_path_join, random_hex, TimeoutError
 
 # classes for config
 from .auth import Authenticator, PAMAuthenticator
@@ -57,10 +56,13 @@ aliases = {
 }
 
 flags = {
-    'debug': ({'Application' : {'log_level' : logging.DEBUG}},
+    'debug': ({'Application' : {'log_level': logging.DEBUG}},
         "set log level to logging.DEBUG (maximize logging output)"),
-    'generate-config': ({'JupyterHubApp': {'generate_config' : True}},
+    'generate-config': ({'JupyterHubApp': {'generate_config': True}},
-        "generate default config file")
+        "generate default config file"),
+    'no-db': ({'JupyterHubApp': {'db_url': 'sqlite:///:memory:'}},
+        "disable persisting state database to disk"
+    ),
 }
 
 
@@ -150,10 +152,13 @@ class JupyterHubApp(Application):
         """
     )
     proxy_auth_token = Unicode(config=True,
-        help="The Proxy Auth token"
+        help="""The Proxy Auth token.
+
+        Loaded from the CONFIGPROXY_AUTH_TOKEN env variable by default.
+        """
     )
     def _proxy_auth_token_default(self):
-        return orm.new_token()
+        return os.environ.get('CONFIGPROXY_AUTH_TOKEN', orm.new_token())
     
     proxy_api_ip = Unicode('localhost', config=True,
         help="The ip for the proxy API handlers"
@@ -190,11 +195,16 @@ class JupyterHubApp(Application):
         if newnew != new:
             self.hub_prefix = newnew
     
-    cookie_secret = Bytes(config=True)
+    cookie_secret = Unicode(config=True,
+        help="""The cookie secret to use to encrypt cookies.
+
+        Loaded from the JPY_COOKIE_SECRET env variable by default.
+        """
+    )
     def _cookie_secret_default(self):
-        return b'secret!'
+        return os.environ.get('JPY_COOKIE_SECRET', random_hex(64))
     
-    authenticator = DottedObjectName("jupyterhub.auth.PAMAuthenticator", config=True,
+    authenticator_class = Type("jupyterhub.auth.PAMAuthenticator", config=True,
         help="""Class for authenticating users.
         
         This should be a class with the following form:
@@ -208,16 +218,38 @@ class JupyterHubApp(Application):
           and `data` is the POST form data from the login page.
         """
     )
+    authenticator = Instance(Authenticator)
+    def _authenticator_default(self):
+        return self.authenticator_class(config=self.config)
+
     # class for spawning single-user servers
-    spawner_class = DottedObjectName("jupyterhub.spawner.LocalProcessSpawner", config=True,
+    spawner_class = Type("jupyterhub.spawner.LocalProcessSpawner", config=True,
         help="""The class to use for spawning single-user servers.
         
         Should be a subclass of Spawner.
         """
     )
     
-    db_url = Unicode('sqlite:///:memory:', config=True)
+    db_url = Unicode('sqlite:///jupyterhub.sqlite', config=True,
-    debug_db = Bool(False)
+        help="url for the database. e.g. `sqlite:///jupyterhub.sqlite`"
+    )
+    def _db_url_changed(self, name, old, new):
+        if '://' not in new:
+            # assume sqlite, if given as a plain filename
+            self.db_url = 'sqlite:///%s' % new
+
+    db_kwargs = Dict(config=True,
+        help="""Include any kwargs to pass to the database connection.
+        See sqlalchemy.create_engine for details.
+        """
+    )
+
+    reset_db = Bool(False, config=True,
+        help="Purge and reset the database."
+    )
+    debug_db = Bool(False, config=True,
+        help="log all database transactions. This has A LOT of output"
+    )
     db = Any()
     
     admin_users = Set({getpass.getuser()}, config=True,
@@ -288,7 +320,6 @@ class JupyterHubApp(Application):
 
         self.handlers = self.add_url_prefix(self.hub_prefix, h)
 
-
         # some extra handlers, outside hub_prefix
         self.handlers.extend([
             (r"%s" % self.hub_prefix.rstrip('/'), web.RedirectHandler,
@@ -302,48 +333,139 @@ class JupyterHubApp(Application):
         ])
     
     def init_db(self):
-        # TODO: load state from db for resume
+        """Create the database connection"""
-        # TODO: if not resuming, clear existing db contents
+        self.db = orm.new_session(self.db_url, reset=self.reset_db, echo=self.debug_db,
-        self.db = orm.new_session(self.db_url, echo=self.debug_db)
+            **self.db_kwargs
-        for name in self.admin_users:
+        )
-            user = orm.User(name=name, admin=True)
-            self.db.add(user)
-        self.db.commit()
     
     def init_hub(self):
         """Load the Hub config into the database"""
-        self.hub = orm.Hub(
+        self.hub = self.db.query(orm.Hub).first()
-            server=orm.Server(
+        if self.hub is None:
-                ip=self.hub_ip,
+            self.hub = orm.Hub(
-                port=self.hub_port,
+                server=orm.Server(
-                base_url=self.hub_prefix,
+                    ip=self.hub_ip,
-                cookie_secret=self.cookie_secret,
+                    port=self.hub_port,
-                cookie_name='jupyter-hub-token',
+                    base_url=self.hub_prefix,
+                    cookie_secret=self.cookie_secret,
+                    cookie_name='jupyter-hub-token',
+                )
             )
-        )
+            self.db.add(self.hub)
-        self.db.add(self.hub)
+        else:
+            server = self.hub.server
+            server.ip = self.hub_ip
+            server.port = self.hub_port
+            server.base_url = self.hub_prefix
+
         self.db.commit()
     
+    def init_users(self):
+        """Load users into and from the database"""
+        db = self.db
+
+        for name in self.admin_users:
+            # ensure anyone specified as admin in config is admin in db
+            user = orm.User.find(db, name)
+            if user is None:
+                user = orm.User(name=name, admin=True)
+                db.add(user)
+            else:
+                user.admin = True
+
+        # the admin_users config variable will never be used after this point.
+        # only the database values will be referenced.
+
+        whitelist = self.authenticator.whitelist
+
+        if not whitelist:
+            self.log.info("Not using whitelist. Any authenticated user will be allowed.")
+
+        # add whitelisted users to the db
+        for name in whitelist:
+            user = orm.User.find(db, name)
+            if user is None:
+                user = orm.User(name=name)
+                db.add(user)
+
+        if whitelist:
+            # fill the whitelist with any users loaded from the db,
+            # so we are consistent in both directions.
+            # This lets whitelist be used to set up initial list,
+            # but changes to the whitelist can occur in the database,
+            # and persist across sessions.
+            for user in db.query(orm.User):
+                whitelist.add(user.name)
+
+        # The whitelist set and the users in the db are now the same.
+        # From this point on, any user changes should be done simultaneously
+        # to the whitelist set and user db, unless the whitelist is empty (all users allowed).
+
+        db.commit()
+
+        # load any still-active spawners from JSON
+        run_sync = IOLoop().run_sync
+
+        user_summaries = ['']
+        def _user_summary(user):
+            parts = ['{0: >8}'.format(user.name)]
+            if user.admin:
+                parts.append('admin')
+            if user.server:
+                parts.append('running at %s' % user.server)
+            return ' '.join(parts)
+
+        for user in db.query(orm.User):
+            if not user.state:
+                user_summaries.append(_user_summary(user))
+                continue
+            self.log.debug("Loading state for %s from db", user.name)
+            spawner = self.spawner_class.fromJSON(user.state, user=user, hub=self.hub, config=self.config)
+            status = run_sync(spawner.poll)
+            if status is None:
+                self.log.info("User %s still running", user.name)
+                user.spawner = spawner
+            else:
+                self.log.warn("Failed to load state for %s, assuming server is not running.", user.name)
+                # not running, state is invalid
+                user.state = {}
+                user.server = None
+
+            user_summaries.append(_user_summary(user))
+
+        self.log.debug("Loaded users: %s", '\n'.join(user_summaries))
+        db.commit()
+
     def init_proxy(self):
         """Load the Proxy config into the database"""
-        self.proxy = orm.Proxy(
+        self.proxy = self.db.query(orm.Proxy).first()
-            public_server=orm.Server(
+        if self.proxy is None:
-                ip=self.ip,
+            self.proxy = orm.Proxy(
-                port=self.port,
+                public_server=orm.Server(),
-            ),
+                api_server=orm.Server(),
-            api_server=orm.Server(
+                auth_token = self.proxy_auth_token,
-                ip=self.proxy_api_ip,
+            )
-                port=self.proxy_api_port,
+            self.db.add(self.proxy)
-                base_url='/api/routes/'
+            self.db.commit()
-            ),
+        self.proxy.log = self.log
-            auth_token = orm.new_token(),
+        self.proxy.public_server.ip = self.ip
-        )
+        self.proxy.public_server.port = self.port
-        self.db.add(self.proxy)
+        self.proxy.api_server.ip = self.proxy_api_ip
+        self.proxy.api_server.port = self.proxy_api_port
+        self.proxy.api_server.base_url = '/api/routes/'
+        if self.proxy.auth_token is None:
+            self.proxy.auth_token = self.proxy_auth_token
         self.db.commit()
     
     @gen.coroutine
     def start_proxy(self):
         """Actually start the configurable-http-proxy"""
+        if self.proxy.public_server.is_up() and \
+              self.proxy.api_server.is_up():
+           self.log.warn("Proxy already running at: %s", self.proxy.public_server.url)
+           self.proxy_process = None
+           return
+
         env = os.environ.copy()
         env['CONFIGPROXY_AUTH_TOKEN'] = self.proxy.auth_token
         cmd = [self.proxy_cmd,
@@ -385,7 +507,9 @@ class JupyterHubApp(Application):
     def check_proxy(self):
         if self.proxy_process.poll() is None:
             return
-        self.log.error("Proxy stopped with exit code %i", self.proxy_process.poll())
+        self.log.error("Proxy stopped with exit code %r",
+            'unknown' if self.proxy_process is None else self.proxy_process.poll()
+        )
         yield self.start_proxy()
         self.log.info("Setting up routes on new proxy")
         yield self.proxy.add_all_users()
@@ -407,10 +531,10 @@ class JupyterHubApp(Application):
             proxy=self.proxy,
             hub=self.hub,
             admin_users=self.admin_users,
-            authenticator=import_item(self.authenticator)(config=self.config),
+            authenticator=self.authenticator,
-            spawner_class=import_item(self.spawner_class),
+            spawner_class=self.spawner_class,
             base_url=base_url,
-            cookie_secret=self.cookie_secret,
+            cookie_secret=self.hub.server.cookie_secret,
             login_url=url_path_join(self.hub.server.base_url, 'login'),
             static_path=os.path.join(self.data_files_path, 'static'),
             static_url_prefix=url_path_join(self.hub.server.base_url, 'static/'),
@@ -438,12 +562,13 @@ class JupyterHubApp(Application):
         if self.generate_config:
             return
         self.load_config_file(self.config_file)
-        self.write_pid_file()
         self.init_logging()
+        self.write_pid_file()
         self.init_ports()
         self.init_db()
         self.init_hub()
         self.init_proxy()
+        self.init_users()
         self.init_handlers()
         self.init_tornado_settings()
         self.init_tornado_application()
@@ -456,15 +581,24 @@ class JupyterHubApp(Application):
         futures = []
         for user in self.db.query(orm.User):
             if user.spawner is not None:
-                futures.append(user.spawner.stop())
+                futures.append(user.stop())
         
         # clean up proxy while SUS are shutting down
-        self.log.info("Cleaning up proxy[%i]..." % self.proxy_process.pid)
+        if self.proxy_process and self.proxy_process.poll() is None:
-        self.proxy_process.terminate()
+            self.log.info("Cleaning up proxy[%i]...", self.proxy_process.pid)
+            try:
+                self.proxy_process.terminate()
+            except Exception as e:
+                self.log.error("Failed to terminate proxy process: %s", e)
         
         # wait for the requests to stop finish:
         for f in futures:
-            yield f
+            try:
+                yield f
+            except Exception as e:
+                self.log.error("Failed to stop user: %s", e)
+
+        self.db.commit()
         
         if self.pid_file and os.path.exists(self.pid_file):
             self.log.info("Cleaning up PID file %s", self.pid_file)
@@ -474,6 +608,7 @@ class JupyterHubApp(Application):
         self.log.info("...done")
     
     def write_config_file(self):
+        """Write our default config to a .py config file"""
         if os.path.exists(self.config_file) and not self.answer_yes:
             answer = ''
             def ask():
@@ -509,9 +644,14 @@ class JupyterHubApp(Application):
             return
         
         loop = IOLoop.current()
+        loop.add_callback(self.proxy.add_all_users)
         
-        pc = PeriodicCallback(self.check_proxy, self.proxy_check_interval)
+        if self.proxy_process:
-        pc.start()
+            # only check / restart the proxy if we started it in the first place.
+            # this means a restarted Hub cannot restart a Proxy that its
+            # predecessor started.
+            pc = PeriodicCallback(self.check_proxy, self.proxy_check_interval)
+            pc.start()
         
         # start the webserver
         http_server = tornado.httpserver.HTTPServer(self.tornado_application)

jupyterhub/handlers/base.py

@@ -73,7 +73,7 @@ class BaseHandler(RequestHandler):
         if not match:
             return None
         token = match.group(1)
-        orm_token = self.db.query(orm.APIToken).filter(orm.APIToken.token == token).first()
+        orm_token = orm.APIToken.find(self.db, token)
         if orm_token is None:
             return None
         else:
@@ -83,8 +83,7 @@ class BaseHandler(RequestHandler):
         """get_current_user from a cookie token"""
         token = self.get_cookie(self.hub.server.cookie_name, None)
         if token:
-            cookie_token = self.db.query(orm.CookieToken).filter(
+            cookie_token = orm.CookieToken.find(self.db, token)
-                orm.CookieToken.token==token).first()
             if cookie_token:
                 return cookie_token.user
             else:
@@ -103,7 +102,7 @@ class BaseHandler(RequestHandler):
         
         return None if no such user
         """
-        return self.db.query(orm.User).filter(orm.User.name==name).first()
+        return orm.User.find(self.db, name)
 
     def user_from_username(self, username):
         """Get ORM User for username"""

jupyterhub/handlers/login.py

@@ -51,7 +51,12 @@ class LoginHandler(BaseHandler):
         authorized = yield self.authenticate(data)
         if authorized:
             user = self.user_from_username(username)
-            yield self.spawn_single_user(user)
+            already_running = False
+            if user.spawner:
+                status = yield user.spawner.poll()
+                already_running = (status == None)
+            if not already_running:
+                yield self.spawn_single_user(user)
             self.set_login_cookie(user)
             next_url = self.get_argument('next', default='') or self.hub.server.base_url
             self.redirect(next_url)

jupyterhub/orm.py

@@ -9,6 +9,7 @@ import socket
 import uuid
 
 from tornado import gen
+from tornado.log import app_log
 from tornado.httpclient import HTTPRequest, AsyncHTTPClient, HTTPError
 
 from sqlalchemy.types import TypeDecorator, VARCHAR
@@ -71,7 +72,7 @@ class Server(Base):
     ip = Column(Unicode, default=u'localhost')
     port = Column(Integer, default=random_port)
     base_url = Column(Unicode, default=u'/')
-    cookie_secret = Column(Binary, default=b'secret')
+    cookie_secret = Column(Unicode, default=u'')
     cookie_name = Column(Unicode, default=u'cookie')
     
     def __repr__(self):
@@ -103,12 +104,11 @@ class Server(Base):
             socket.create_connection((self.ip or 'localhost', self.port))
         except socket.error as e:
             if e.errno == errno.ECONNREFUSED:
-                return True
+                return False
             else:
                 raise
         else:
             return True
-        
 
 
 class Proxy(Base):
@@ -124,6 +124,7 @@ class Proxy(Base):
     public_server = relationship(Server, primaryjoin=_public_server_id == Server.id)
     _api_server_id = Column(Integer, ForeignKey('servers.id'))
     api_server = relationship(Server, primaryjoin=_api_server_id == Server.id)
+    log = app_log
     
     def __repr__(self):
         if self.public_server:
@@ -136,6 +137,9 @@ class Proxy(Base):
     @gen.coroutine
     def add_user(self, user, client=None):
         """Add a user's server to the proxy table."""
+        self.log.info("Adding user %s to proxy %s => %s",
+            user.name, user.server.base_url, user.server.host,
+        )
         client = client or AsyncHTTPClient()
         
         req = HTTPRequest(url_path_join(
@@ -155,6 +159,7 @@ class Proxy(Base):
     @gen.coroutine
     def delete_user(self, user, client=None):
         """Remove a user's server to the proxy table."""
+        self.log.info("Removing user %s from proxy", user.name)
         client = client or AsyncHTTPClient()
         req = HTTPRequest(url_path_join(
                 self.api_server.url,
@@ -262,6 +267,14 @@ class User(Base):
         """Return a new cookie token"""
         return self._new_token(CookieToken)
 
+    @classmethod
+    def find(cls, db, name):
+        """Find a user by name.
+
+        Returns None if not found.
+        """
+        return db.query(cls).filter(cls.name==name).first()
+
     @gen.coroutine
     def spawn(self, spawner_class, base_url='/', hub=None, config=None):
         db = inspect(self).session
@@ -321,6 +334,14 @@ class Token(object):
             u=self.user.name,
         )
 
+    @classmethod
+    def find(cls, db, token):
+        """Find a token object by value.
+
+        Returns None if not found.
+        """
+        return db.query(cls).filter(cls.token==token).first()
+
 
 class APIToken(Token, Base):
     """An API token"""
@@ -332,13 +353,15 @@ class CookieToken(Token, Base):
     __tablename__ = 'cookie_tokens'
 
 
-def new_session(url="sqlite:///:memory:", **kwargs):
+def new_session(url="sqlite:///:memory:", reset=False, **kwargs):
     """Create a new session at url"""
     kwargs.setdefault('connect_args', {'check_same_thread': False})
     kwargs.setdefault('poolclass', StaticPool)
     engine = create_engine(url, **kwargs)
     Session = sessionmaker(bind=engine)
     session = Session()
+    if reset:
+        Base.metadata.drop_all(engine)
     Base.metadata.create_all(engine)
     return session
 

jupyterhub/spawner.py

@@ -60,7 +60,7 @@ class Spawner(LoggingConfigurable):
         env = os.environ.copy()
         for key in ['HOME', 'USER', 'USERNAME', 'LOGNAME', 'LNAME']:
             env.pop(key, None)
-        self._env_key(env, 'COOKIE_SECRET', self.user.server.cookie_secret.decode('ascii'))
+        self._env_key(env, 'COOKIE_SECRET', self.user.server.cookie_secret)
         self._env_key(env, 'API_TOKEN', self.api_token)
         return env
     
@@ -103,7 +103,7 @@ class Spawner(LoggingConfigurable):
         state: dict
              a JSONable dict of state
         """
-        return {}
+        return dict(api_token=self.api_token)
     
     def get_args(self):
         """Return the arguments to be passed after self.cmd"""
@@ -211,10 +211,13 @@ class LocalProcessSpawner(Spawner):
             raise ValueError("This should be impossible")
     
     def load_state(self, state):
+        super(LocalProcessSpawner, self).load_state(state)
         self.pid = state['pid']
     
     def get_state(self):
-        return dict(pid=self.pid)
+        state = super(LocalProcessSpawner, self).get_state()
+        state['pid'] = self.pid
+        return state
     
     def sudo_cmd(self, user):
         return ['sudo', '-u', user.name] + self.sudo_args

jupyterhub/tests/mocking.py

@@ -56,11 +56,14 @@ class MockHubApp(JupyterHubApp):
     def _ip_default(self):
         return 'localhost'
     
-    def _authenticator_default(self):
+    def _db_url_default(self):
-        return '%s.%s' % (__name__, 'MockPAMAuthenticator')
+        return 'sqlite:///:memory:'
+
+    def _authenticator_class_default(self):
+        return MockPAMAuthenticator
     
     def _spawner_class_default(self):
-        return '%s.%s' % (__name__, 'MockSpawner')
+        return MockSpawner
     
     def _admin_users_default(self):
         return {'admin'}

jupyterhub/tests/test_orm.py

@@ -21,7 +21,7 @@ def test_server(db):
     assert server.proto == 'http'
     assert isinstance(server.port, int)
     assert isinstance(server.cookie_name, unicode)
-    assert isinstance(server.cookie_secret, bytes)
+    assert isinstance(server.cookie_secret, unicode)
     assert server.url == 'http://localhost:%i/' % server.port
 
 
@@ -71,6 +71,11 @@ def test_user(db):
     assert user.server.ip == u'localhost'
     assert user.state == {'pid': 4234}
 
+    found = orm.User.find(db, u'kaylee')
+    assert found.name == user.name
+    found = orm.User.find(db, u'badger')
+    assert found is None
+
 
 def test_tokens(db):
     user = orm.User(name=u'inara')
@@ -87,3 +92,7 @@ def test_tokens(db):
     assert len(user.api_tokens) == 1
     assert len(user.cookie_tokens) == 3
     
+    found = orm.CookieToken.find(db, token=token.token)
+    assert found.token == token.token
+    found = orm.APIToken.find(db, token.token)
+    assert found is None

jupyterhub/utils.py

@@ -3,7 +3,9 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+import binascii
 import errno
+import os
 import socket
 from tornado import web, gen, ioloop
 from tornado.log import app_log
@@ -11,7 +13,8 @@ from tornado.log import app_log
 from IPython.html.utils import url_path_join
 
 try:
-    TimeoutError
+    # make TimeoutError importable on Python >= 3.3
+    TimeoutError = TimeoutError
 except NameError:
     # python < 3.3
     class TimeoutError(Exception):
@@ -25,6 +28,12 @@ def random_port():
     sock.close()
     return port
 
+def random_hex(nbytes):
+    """Return nbytes random bytes as a unicode hex string
+
+    It will have length nbytes * 2
+    """
+    return binascii.hexlify(os.urandom(nbytes)).decode('ascii')
 
 @gen.coroutine
 def wait_for_server(ip, port, timeout=10):