Merge pull request #117 from minrk/cli-token

allow creating new API tokens on the command line

examples/cull-idle/cull_idle_servers.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python
+"""script to monitor and cull idle single-user servers
+
+Caveats:
+
+last_activity is not updated with high frequency,
+so cull timeout should be greater than the sum of:
+
+- single-user websocket ping interval (default: 30s)
+- JupyterHub.last_activity_interval (default: 5 minutes)
+
+Generate an API token and store it in `JPY_API_TOKEN`:
+
+    export JPY_API_TOKEN=`jupyterhub token`
+    python cull_idle_servers.py [--timeout=900] [--url=http://127.0.0.1:8081/hub]
+"""
+
+import datetime
+import json
+import os
+
+from dateutil.parser import parse as parse_date
+
+from tornado.gen import coroutine
+from tornado.log import app_log
+from tornado.httpclient import AsyncHTTPClient, HTTPRequest
+from tornado.ioloop import IOLoop, PeriodicCallback
+from tornado.options import define, options, parse_command_line
+
+
+@coroutine
+def cull_idle(url, api_token, timeout):
+    """cull idle single-user servers"""
+    auth_header = {
+            'Authorization': 'token %s' % api_token
+        }
+    req = HTTPRequest(url=url + '/api/users',
+        headers=auth_header,
+    )
+    now = datetime.datetime.utcnow()
+    cull_limit = now - datetime.timedelta(seconds=timeout)
+    client = AsyncHTTPClient()
+    resp = yield client.fetch(req)
+    users = json.loads(resp.body.decode('utf8', 'replace'))
+    futures = []
+    for user in users:
+        last_activity = parse_date(user['last_activity'])
+        if user['server'] and last_activity < cull_limit:
+            app_log.info("Culling %s (inactive since %s)", user['name'], last_activity)
+            req = HTTPRequest(url=url + '/api/users/%s/server' % user['name'],
+                method='DELETE',
+                headers=auth_header,
+            )
+            futures.append((user['name'], client.fetch(req)))
+        elif user['server'] and last_activity > cull_limit:
+            app_log.debug("Not culling %s (active since %s)", user['name'], last_activity)
+    
+    for (name, f) in futures:
+        yield f
+        app_log.debug("Finished culling %s", name)
+
+if __name__ == '__main__':
+    define('url', default='http://127.0.0.1:8081/hub', help="The JupyterHub API URL")
+    define('timeout', default=600, help="The idle timeout (in seconds)")
+    define('cull_every', default=0, help="The interval (in seconds) for checking for idle servers to cull")
+    
+    parse_command_line()
+    if not options.cull_every:
+        options.cull_every = options.timeout // 2
+    
+    api_token = os.environ['JPY_API_TOKEN']
+    
+    loop = IOLoop.current()
+    cull = lambda : cull_idle(options.url, api_token, options.timeout)
+    # run once before scheduling periodic call
+    loop.run_sync(cull)
+    # schedule periodic cull
+    pc = PeriodicCallback(cull, 1e3 * options.cull_every)
+    pc.start()
+    try:
+        loop.start()
+    except KeyboardInterrupt:
+        pass
+    

jupyterhub/apihandlers/users.py

@@ -18,6 +18,7 @@ class BaseUserHandler(APIHandler):
             'name': user.name,
             'admin': user.admin,
             'server': user.server.base_url if user.server else None,
+            'last_activity': user.last_activity.isoformat(),
         }
     
     _model_types = {

jupyterhub/app.py

@@ -79,6 +79,44 @@ flags = {
 
 SECRET_BYTES = 2048 # the number of bytes to use when generating new secrets
 
+class NewToken(Application):
+    """Generate and print a new API token"""
+    name = 'jupyterhub-token'
+    description = """Generate and return new API token for a user.
+    
+    Usage:
+    
+        jupyterhub token [username]
+    """
+    
+    examples = """
+        $> jupyterhub token kaylee
+        ab01cd23ef45
+    """
+    
+    name = Unicode(getuser())
+    
+    def parse_command_line(self, argv=None):
+        super().parse_command_line(argv=argv)
+        if not self.extra_args:
+            return
+        if len(self.extra_args) > 1:
+            print("Must specify exactly one username", file=sys.stderr)
+            self.exit(1)
+        self.name = self.extra_args[0]
+    
+    def start(self):
+        hub = JupyterHub(parent=self)
+        hub.init_db()
+        hub.init_users()
+        user = orm.User.find(hub.db, self.name)
+        if user is None:
+            print("No such user: %s" % self.name)
+            self.exit(1)
+        token = user.new_api_token()
+        print(token)
+
+
 class JupyterHub(Application):
     """An Application for starting a Multi-User Jupyter Notebook server."""
     name = 'jupyterhub'
@@ -104,6 +142,10 @@ class JupyterHub(Application):
     aliases = Dict(aliases)
     flags = Dict(flags)
     
+    subcommands = {
+        'token': (NewToken, "Generate an API token for a user")
+    }
+    
     classes = List([
         Spawner,
         LocalProcessSpawner,
@@ -125,7 +167,7 @@ class JupyterHub(Application):
         Useful for daemonizing jupyterhub.
         """
     )
-    last_activity_interval = Integer(600, config=True,
+    last_activity_interval = Integer(300, config=True,
         help="Interval (in seconds) at which to update last-activity timestamps."
     )
     proxy_check_interval = Integer(30, config=True,
@@ -695,7 +737,7 @@ class JupyterHub(Application):
     @catch_config_error
     def initialize(self, *args, **kwargs):
         super().initialize(*args, **kwargs)
-        if self.generate_config:
+        if self.generate_config or self.subapp:
             return
         self.load_config_file(self.config_file)
         self.init_logging()
@@ -800,6 +842,9 @@ class JupyterHub(Application):
 
     def start(self):
         """Start the whole thing"""
+        if self.subapp:
+            return self.subapp.start()
+        
         if self.generate_config:
             self.write_config_file()
             return

jupyterhub/handlers/base.py

@@ -79,9 +79,7 @@ class BaseHandler(RequestHandler):
         if orm_token is None:
             return None
         else:
-            user = orm_token.user
-            user.last_activity = datetime.utcnow()
-            return user
+            return orm_token.user
     
     def _user_for_cookie(self, cookie_name, cookie_value=None):
         """Get the User for a given cookie, if there is one"""

jupyterhub/tests/test_api.py

@@ -95,7 +95,11 @@ def test_get_users(app):
     db = app.db
     r = api_request(app, 'users')
     assert r.status_code == 200
-    assert sorted(r.json(), key=lambda d: d['name']) == [
+    
+    users = sorted(r.json(), key=lambda d: d['name'])
+    for u in users:
+        u.pop('last_activity')
+    assert users == [
         {
             'name': 'admin',
             'admin': True,

jupyterhub/tests/test_app.py

@@ -1,15 +1,24 @@
 """Test the JupyterHub entry point"""
 
 import os
+import re
 import sys
+from getpass import getuser
 from subprocess import check_output
-from tempfile import NamedTemporaryFile
+from tempfile import NamedTemporaryFile, TemporaryDirectory
 
 def test_help_all():
     out = check_output([sys.executable, '-m', 'jupyterhub', '--help-all']).decode('utf8', 'replace')
     assert '--ip' in out
     assert '--JupyterHub.ip' in out
 
+def test_token_app():
+    cmd = [sys.executable, '-m', 'jupyterhub', 'token']
+    out = check_output(cmd + ['--help-all']).decode('utf8', 'replace')
+    with TemporaryDirectory() as td:
+        out = check_output(cmd + [getuser()], cwd=td).decode('utf8', 'replace').strip()
+    assert re.match(r'^[a-z0-9]+$', out)
+
 def test_generate_config():
     with NamedTemporaryFile(prefix='jupyterhub_config', suffix='.py') as tf:
         cfg_file = tf.name