drop-token

#!/usr/bin/env python3
"""Drop a token from a JupyterHub database

Usage:

    ./drop-token abcdef [abc123] ...
"""

from jupyterhub import orm

class KeyNotFound(Exception):
    pass

def remove_token(*tokens, db_url='sqlite:///jupyterhub.sqlite'):
    """Remove a given token from the database"""

    db = orm.new_session_factory(url=db_url)()
    for token in tokens:
        orm_token = orm.APIToken.find(db, token)
        if orm_token is None:
            raise KeyNotFound("No token matching %s" % token)
        else:
            print("Removing %s" % orm_token)
            db.delete(orm_token)
    db.commit()
    
if __name__ == '__main__':
    import argparse
    import sys
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("token", nargs='+', help="the token(s) to remove")
    parser.add_argument("--db", default='sqlite:///jupyterhub.sqlite',
        help="The db url. Default: jupyterhub.sqlite in the CWD")
    opts = parser.parse_args()
    try:
        remove_token(*opts.tokens, db_url=opts.db)
    except KeyNotFound as e:
        sys.exit(str(e))