add user groups

jupyterhub/apihandlers/base.py

@@ -90,6 +90,7 @@ class APIHandler(BaseHandler):
         model = {
             'name': user.name,
             'admin': user.admin,
+            'groups': [ g.name for g in user.groups ],
             'server': user.url if user.running else None,
             'pending': None,
             'last_activity': user.last_activity.isoformat(),

jupyterhub/orm.py

@@ -20,7 +20,7 @@ from sqlalchemy.ext.declarative import declarative_base, declared_attr
 from sqlalchemy.orm import sessionmaker, relationship
 from sqlalchemy.pool import StaticPool
 from sqlalchemy.sql.expression import bindparam
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, Table
 
 from .utils import (
     random_port, url_path_join, wait_for_server, wait_for_http_server,
@@ -258,6 +258,32 @@ class Hub(Base):
             return "<%s [unconfigured]>" % self.__class__.__name__
 
 
+# user:group many:many mapping table
+user_group_map = Table('user_group_map', Base.metadata,
+    Column('user_id', ForeignKey('users.id'), primary_key=True),
+    Column('group_id', ForeignKey('groups.id'), primary_key=True),
+)
+
+class Group(Base):
+    """User Groups"""
+    __tablename__ = 'groups'
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    name = Column(Unicode(1023), unique=True)
+    users = relationship('User', secondary='user_group_map', back_populates='groups')
+    
+    def __repr__(self):
+        return "<%s %s (%i users)>" % (
+            self.__class__.__name__, self.name, len(self.users)
+        )
+    @classmethod
+    def find(cls, db, name):
+        """Find a group by name.
+
+        Returns None if not found.
+        """
+        return db.query(cls).filter(cls.name==name).first()
+
+
 class User(Base):
     """The User table
 
@@ -290,6 +316,8 @@ class User(Base):
     state = Column(JSONDict)
     # Authenticators can store their state here:
     auth_state = Column(JSONDict)
+    # group mapping
+    groups = relationship('Group', secondary='user_group_map', back_populates='users')
 
     other_user_cookies = set([])
 

jupyterhub/tests/test_api.py

@@ -21,14 +21,13 @@ def check_db_locks(func):
     Decorator for test functions that verifies no locks are held on the
     application's database upon exit by creating and dropping a dummy table.
 
-    Relies on an instance of JupyterhubApp being the first argument to the
+    Relies on an instance of JupyterHubApp being the first argument to the
     decorated function.
     """
 
-    def new_func(*args, **kwargs):
+    def new_func(app, *args, **kwargs):
-        retval = func(*args, **kwargs)
+        retval = func(app, *args, **kwargs)
 
-        app = args[0]
         temp_session = app.session_factory()
         temp_session.execute('CREATE TABLE dummy (foo INT)')
         temp_session.execute('DROP TABLE dummy')
@@ -159,12 +158,14 @@ def test_get_users(app):
     assert users == [
         {
             'name': 'admin',
+            'groups': [],
             'admin': True,
             'server': None,
             'pending': None,
         },
         {
             'name': 'user',
+            'groups': [],
             'admin': False,
             'server': None,
             'pending': None,
@@ -195,6 +196,7 @@ def test_get_user(app):
     user.pop('last_activity')
     assert user == {
         'name': name,
+        'groups': [],
         'admin': False,
         'server': None,
         'pending': None,

jupyterhub/tests/test_orm.py

@@ -124,3 +124,17 @@ def test_spawn_fails(db, io_loop):
     assert user.server is None
     assert not user.running
 
+
+def test_groups(db):
+    user = orm.User(name='aeofel')
+    db.add(user)
+    
+    group = orm.Group(name='lives')
+    db.add(group)
+    db.commit()
+    assert group.users == []
+    assert user.groups == []
+    group.users.append(user)
+    db.commit()
+    assert group.users == [user]
+    assert user.groups == [group]