From 4e7e586cb9396dc954cc0209804584a6179733a2 Mon Sep 17 00:00:00 2001
From: Min RK <benjaminrk@gmail.com>
Date: Fri, 13 Apr 2018 14:15:03 +0200
Subject: [PATCH] add `python -m jupyterhub.dbutil shell`

allows opening an IPython shell with a connection to your database

alembic moved from `python -m jupyterhub.dbutil` to `python -m jupyterhub.dbutil alembic` subcommand
---
 jupyterhub/dbutil.py | 46 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 42 insertions(+), 4 deletions(-)

diff --git a/jupyterhub/dbutil.py b/jupyterhub/dbutil.py
index ee608524..98a338fe 100644
--- a/jupyterhub/dbutil.py
+++ b/jupyterhub/dbutil.py
@@ -130,13 +130,51 @@ def upgrade_if_needed(db_url, backup=True, log=None):
     upgrade(db_url)
 
 
-def _alembic(*args):
+def shell(args=None):
+    """Start an IPython shell hooked up to the jupyerhub database"""
+    from .app import JupyterHub
+    hub = JupyterHub()
+    hub.load_config_file(hub.config_file)
+    db_url = hub.db_url
+    db = orm.new_session_factory(db_url, **hub.db_kwargs)()
+    ns = {
+        'db': db,
+        'db_url': db_url,
+        'orm': orm,
+    }
+
+    import IPython
+    IPython.start_ipython(args, user_ns=ns)
+
+
+def _alembic(args):
     """Run an alembic command with a temporary alembic.ini"""
-    with _temp_alembic_ini('sqlite:///jupyterhub.sqlite') as alembic_ini:
+    from .app import JupyterHub
+    hub = JupyterHub()
+    hub.load_config_file(hub.config_file)
+    db_url = hub.db_url
+    with _temp_alembic_ini(db_url) as alembic_ini:
         check_call(
-            ['alembic', '-c', alembic_ini] + list(args)
+            ['alembic', '-c', alembic_ini] + args
         )
 
 
+def main(args=None):
+    if args is None:
+        args = sys.argv[1:]
+    # dumb option parsing, since we want to pass things through
+    # to subcommands
+    choices = ['shell', 'alembic']
+    if not args or args[0] not in choices:
+        print("Select a command from: %s" % ', '.join(choices))
+        return 1
+    cmd, args = args[0], args[1:]
+
+    if cmd == 'shell':
+        shell(args)
+    elif cmd == 'alembic':
+        _alembic(args)
+
+
 if __name__ == '__main__':
-    _alembic(*sys.argv[1:])
+    sys.exit(main())