Merge pull request #907 from barrachri/adding_multi_server_features_orm

Allow users to have multiple servers (in db)

jupyterhub/app.py

@@ -1095,7 +1095,10 @@ class JupyterHub(Application):
                 # if user.server is defined.
                 log = self.log.warning if user.server else self.log.debug
                 log("%s not running.", user.name)
-                user.server = None
+                # remove all server or servers entry from db related to the user
+                for server in user.servers:
+                    db.delete(server)
+                db.commit()
 
             user_summaries.append(_user_summary(user))
 

jupyterhub/orm.py

@@ -17,8 +17,10 @@ from sqlalchemy import (
     DateTime,
 )
 from sqlalchemy.ext.declarative import declarative_base, declared_attr
-from sqlalchemy.orm import sessionmaker, relationship
+from sqlalchemy.orm import sessionmaker, relationship, backref
 from sqlalchemy.pool import StaticPool
+from sqlalchemy.schema import Index, UniqueConstraint
+from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.sql.expression import bindparam
 from sqlalchemy import create_engine, Table
 
@@ -68,6 +70,9 @@ class Server(Base):
     base_url = Column(Unicode(255), default='/')
     cookie_name = Column(Unicode(255), default='cookie')
 
+    # added to handle multi-server feature
+    last_activity = Column(DateTime, default=datetime.utcnow)
+
     def __repr__(self):
         return "<Server(%s:%s)>" % (self.ip, self.port)
 
@@ -181,6 +186,8 @@ class Proxy(Base):
             client=client,
         )
 
+    # FIX-ME
+    # we need to add a reference to a specific server
     @gen.coroutine
     def add_user(self, user, client=None):
         """Add a user's server to the proxy table."""
@@ -248,6 +255,8 @@ class Proxy(Base):
         resp = yield self.api_request('', client=client)
         return json.loads(resp.body.decode('utf8', 'replace'))
 
+    # FIX-ME
+    # we need to add a reference to a specific server
     @gen.coroutine
     def check_routes(self, user_dict, service_dict, routes=None):
         """Check that all users are properly routed on the proxy"""
@@ -334,7 +343,6 @@ class Group(Base):
     @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()
@@ -343,8 +351,10 @@ class Group(Base):
 class User(Base):
     """The User table
 
-    Each user has a single server,
-    and multiple tokens used for authorization.
+    Each user can have one or more single user notebook servers.
+
+    Each single user notebook server will have a unique token for authorization.
+    Therefore, a user with multiple notebook servers will have multiple tokens.
 
     API tokens grant access to the Hub's REST API.
     These are used by single-user servers to authenticate requests,
@@ -355,13 +365,17 @@ class User(Base):
 
     A `state` column contains a JSON dict,
     used for restoring state of a Spawner.
+
+
+    `servers` is a list that contains a reference for each of the user's single user notebook servers.
+    The method `server` returns the first entry in the user's `servers` list.
     """
     __tablename__ = 'users'
     id = Column(Integer, primary_key=True, autoincrement=True)
     name = Column(Unicode(1023), unique=True)
-    # should we allow multiple servers per user?
-    _server_id = Column(Integer, ForeignKey('servers.id', ondelete="SET NULL"))
-    server = relationship(Server, primaryjoin=_server_id == Server.id)
+
+    servers = association_proxy("user_to_servers", "server", creator=lambda server: UserServer(server=server))
+
     admin = Column(Boolean, default=False)
     last_activity = Column(DateTime, default=datetime.utcnow)
 
@@ -377,6 +391,16 @@ class User(Base):
 
     other_user_cookies = set([])
 
+    @property
+    def server(self):
+        """Returns the first element of servers.
+        Returns None if the list is empty.
+        """
+        if len(self.servers) == 0:
+            return None
+        else:
+            return self.servers[0]
+
     def __repr__(self):
         if self.server:
             return "<{cls}({name}@{ip}:{port})>".format(
@@ -401,12 +425,40 @@ class User(Base):
     @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()
 
 
+class UserServer(Base):
+    """The UserServer table
+
+    A table storing the One-To-Many relationship between a user and servers.
+    Each user may have one or more servers.
+    A server can have only one (1) user. This condition is maintained by UniqueConstraint.
+    """
+    __tablename__ = 'users_servers'
+
+    _user_id = Column(Integer, ForeignKey('users.id'), primary_key=True)
+    _server_id = Column(Integer, ForeignKey('servers.id'), primary_key=True)
+
+    user = relationship(User, backref=backref('user_to_servers', cascade='all, delete-orphan'))
+    server = relationship(Server, backref=backref('server_to_users', cascade='all, delete-orphan')
+                          )
+
+    __table_args__ = (
+            UniqueConstraint('_server_id'),
+            Index('server_user_index', '_server_id', '_user_id'),)
+
+    def __repr__(self):
+        return "<{cls}({name}@{ip}:{port})>".format(
+            cls=self.__class__.__name__,
+            name=self.user.name,
+            ip=self.server.ip,
+            port=self.server.port,
+        )
+
+
 class Service(Base):
     """A service run with JupyterHub
 
@@ -414,7 +466,6 @@ class Service(Base):
     A service can have API tokens for accessing the Hub's API
 
     It has:
-
     - name
     - admin
     - api tokens
@@ -441,7 +492,6 @@ class Service(Base):
 
     def new_api_token(self, token=None):
         """Create a new API token
-
         If `token` is given, load that token.
         """
         return APIToken.new(token=token, service=self)

jupyterhub/tests/conftest.py

@@ -30,8 +30,8 @@ def db():
         _db = orm.new_session_factory('sqlite:///:memory:', echo=True)()
         user = orm.User(
             name=getuser(),
-            server=orm.Server(),
         )
+        user.servers.append(orm.Server())
         hub = orm.Hub(
             server=orm.Server(),
         )

jupyterhub/tests/test_orm.py

@@ -65,9 +65,10 @@ def test_hub(db):
 
 def test_user(db):
     user = orm.User(name='kaylee',
-        server=orm.Server(),
         state={'pid': 4234},
     )
+    server = orm.Server()
+    user.servers.append(server)
     db.add(user)
     db.commit()
     assert user.name == 'kaylee'

jupyterhub/user.py

@@ -198,12 +198,12 @@ class User(HasTraits):
     def spawn(self, options=None):
         """Start the user's spawner"""
         db = self.db
-        
-        self.server = orm.Server(
+        server = orm.Server(
             cookie_name=self.cookie_name,
             base_url=self.base_url,
         )
-        db.add(self.server)
+        self.servers.append(server)
+        db.add(self)
         db.commit()
 
         api_token = self.new_api_token()
@@ -315,11 +315,10 @@ class User(HasTraits):
             spawner.clear_state()
             self.state = spawner.get_state()
             self.last_activity = datetime.utcnow()
-            # cleanup server entry, API token from defunct server
-            if self.server:
-                # cleanup server entry from db
-                self.db.delete(self.server)
-            self.server = None
+            # Cleanup defunct servers: delete entry and API token for each server
+            for server in self.servers:
+                # remove server entry from db
+                self.db.delete(server)
             if not spawner.will_resume:
                 # find and remove the API token if the spawner isn't
                 # going to re-use it next time
@@ -335,4 +334,3 @@ class User(HasTraits):
                 yield gen.maybe_future(
                     auth.post_spawn_stop(self, spawner)
                 )
-