user sessions

Author: miguelgrinberg

docs/server.rst

@@ -275,6 +275,70 @@ during the broadcast.
     def message(sid, data):
         sio.emit('my reply', data, room='chat_users', skip_sid=sid)
 
+User Sessions
+-------------
+
+The server can maintain application-specific information in a user session
+dedicated to each connected client. Applications can use the user session to
+write any details about the user that need to be preserved throughout the life
+of the connection, such as usernames or user ids.
+
+The ``save_session()`` and ``get_session()`` methods are used to store and
+retrieve information in the user session::
+
+    @sio.on('connect')
+    def on_connect(sid, environ):
+        username = authenticate_user(environ)
+        sio.save_session(sid, {'username': username})
+
+    @sio.on('message')
+    def on_message(sid, data):
+        session = sio.get_session(sid)
+        print('message from ', session['username'])
+
+For the ``asyncio`` server, these methods are coroutines::
+
+
+    @sio.on('connect')
+    async def on_connect(sid, environ):
+        username = authenticate_user(environ)
+        await sio.save_session(sid, {'username': username})
+
+    @sio.on('message')
+    async def on_message(sid, data):
+        session = await sio.get_session(sid)
+        print('message from ', session['username'])
+
+The session can also be manipulated with the `session()` context manager::
+
+    @sio.on('connect')
+    def on_connect(sid, environ):
+        username = authenticate_user(environ)
+        with sio.session(sid) as session:
+            session['username'] = username
+
+    @sio.on('message')
+    def on_message(sid, data):
+        with sio.session(sid) as session:
+            print('message from ', session['username'])
+
+For the ``asyncio`` server, an asynchronous context manager is used::
+
+    @sio.on('connect')
+    def on_connect(sid, environ):
+        username = authenticate_user(environ)
+        async with sio.session(sid) as session:
+            session['username'] = username
+
+    @sio.on('message')
+    def on_message(sid, data):
+        async with sio.session(sid) as session:
+            print('message from ', session['username'])
+
+The ``get_session()``, ``save_session()`` and ``session()`` methods take an
+optional ``namespace`` argument. If this argument isn't provided, the session
+is attached to the default namespace.
+
 Using a Message Queue
 ---------------------
 

socketio/asyncio_namespace.py

@@ -82,6 +82,39 @@ async def close_room(self, room, namespace=None):
         return await self.server.close_room(
             room, namespace=namespace or self.namespace)
 
+    async def get_session(self, sid, namespace=None):
+        """Return the user session for a client.
+
+        The only difference with the :func:`socketio.Server.get_session`
+        method is that when the ``namespace`` argument is not given the
+        namespace associated with the class is used.
+
+        Note: this method is a coroutine.
+        """
+        return await self.server.get_session(
+            sid, namespace=namespace or self.namespace)
+
+    async def save_session(self, sid, session, namespace=None):
+        """Store the user session for a client.
+
+        The only difference with the :func:`socketio.Server.save_session`
+        method is that when the ``namespace`` argument is not given the
+        namespace associated with the class is used.
+
+        Note: this method is a coroutine.
+        """
+        return await self.server.save_session(
+            sid, session, namespace=namespace or self.namespace)
+
+    def session(self, sid, namespace=None):
+        """Return the user session for a client with context manager syntax.
+
+        The only difference with the :func:`socketio.Server.session` method is
+        that when the ``namespace`` argument is not given the namespace
+        associated with the class is used.
+        """
+        return self.server.session(sid, namespace=namespace or self.namespace)
+
     async def disconnect(self, sid, namespace=None):
         """Disconnect a client.
 

socketio/asyncio_server.py

@@ -169,6 +169,73 @@ async def close_room(self, room, namespace=None):
         self.logger.info('room %s is closing [%s]', room, namespace)
         await self.manager.close_room(room, namespace)
 
+    async def get_session(self, sid, namespace=None):
+        """Return the user session for a client.
+
+        :param sid: The session id of the client.
+        :param namespace: The Socket.IO namespace. If this argument is omitted
+                          the default namespace is used.
+
+        The return value is a dictionary. Modifications made to this
+        dictionary are not guaranteed to be preserved. If you want to modify
+        the user session, use the ``session`` context manager instead.
+        """
+        namespace = namespace or '/'
+        eio_session = await self.eio.get_session(sid)
+        return eio_session.setdefault(namespace, {})
+
+    async def save_session(self, sid, session, namespace=None):
+        """Store the user session for a client.
+
+        :param sid: The session id of the client.
+        :param session: The session dictionary.
+        :param namespace: The Socket.IO namespace. If this argument is omitted
+                          the default namespace is used.
+        """
+        namespace = namespace or '/'
+        eio_session = await self.eio.get_session(sid)
+        eio_session[namespace] = session
+
+    def session(self, sid, namespace=None):
+        """Return the user session for a client with context manager syntax.
+
+        :param sid: The session id of the client.
+
+        This is a context manager that returns the user session dictionary for
+        the client. Any changes that are made to this dictionary inside the
+        context manager block are saved back to the session. Example usage::
+
+            @eio.on('connect')
+            def on_connect(sid, environ):
+                username = authenticate_user(environ)
+                if not username:
+                    return False
+                with eio.session(sid) as session:
+                    session['username'] = username
+
+            @eio.on('message')
+            def on_message(sid, msg):
+                async with eio.session(sid) as session:
+                    print('received message from ', session['username'])
+        """
+        class _session_context_manager(object):
+            def __init__(self, server, sid, namespace):
+                self.server = server
+                self.sid = sid
+                self.namespace = namespace
+                self.session = None
+
+            async def __aenter__(self):
+                self.session = await self.server.get_session(
+                    sid, namespace=self.namespace)
+                return self.session
+
+            async def __aexit__(self, *args):
+                await self.server.save_session(sid, self.session,
+                                               namespace=self.namespace)
+
+        return _session_context_manager(self, sid, namespace)
+
     async def disconnect(self, sid, namespace=None):
         """Disconnect a client.
 

socketio/namespace.py

@@ -100,6 +100,35 @@ def rooms(self, sid, namespace=None):
         """
         return self.server.rooms(sid, namespace=namespace or self.namespace)
 
+    def get_session(self, sid, namespace=None):
+        """Return the user session for a client.
+
+        The only difference with the :func:`socketio.Server.get_session`
+        method is that when the ``namespace`` argument is not given the
+        namespace associated with the class is used.
+        """
+        return self.server.get_session(
+            sid, namespace=namespace or self.namespace)
+
+    def save_session(self, sid, session, namespace=None):
+        """Store the user session for a client.
+
+        The only difference with the :func:`socketio.Server.save_session`
+        method is that when the ``namespace`` argument is not given the
+        namespace associated with the class is used.
+        """
+        return self.server.save_session(
+            sid, session, namespace=namespace or self.namespace)
+
+    def session(self, sid, namespace=None):
+        """Return the user session for a client with context manager syntax.
+
+        The only difference with the :func:`socketio.Server.session` method is
+        that when the ``namespace`` argument is not given the namespace
+        associated with the class is used.
+        """
+        return self.server.session(sid, namespace=namespace or self.namespace)
+
     def disconnect(self, sid, namespace=None):
         """Disconnect a client.
 

socketio/server.py

@@ -317,6 +317,74 @@ def rooms(self, sid, namespace=None):
         namespace = namespace or '/'
         return self.manager.get_rooms(sid, namespace)
 
+    def get_session(self, sid, namespace=None):
+        """Return the user session for a client.
+
+        :param sid: The session id of the client.
+        :param namespace: The Socket.IO namespace. If this argument is omitted
+                          the default namespace is used.
+
+        The return value is a dictionary. Modifications made to this
+        dictionary are not guaranteed to be preserved unless
+        ``save_session()`` is called, or when the ``session`` context manager
+        is used.
+        """
+        namespace = namespace or '/'
+        eio_session = self.eio.get_session(sid)
+        return eio_session.setdefault(namespace, {})
+
+    def save_session(self, sid, session, namespace=None):
+        """Store the user session for a client.
+
+        :param sid: The session id of the client.
+        :param session: The session dictionary.
+        :param namespace: The Socket.IO namespace. If this argument is omitted
+                          the default namespace is used.
+        """
+        namespace = namespace or '/'
+        eio_session = self.eio.get_session(sid)
+        eio_session[namespace] = session
+
+    def session(self, sid, namespace=None):
+        """Return the user session for a client with context manager syntax.
+
+        :param sid: The session id of the client.
+
+        This is a context manager that returns the user session dictionary for
+        the client. Any changes that are made to this dictionary inside the
+        context manager block are saved back to the session. Example usage::
+
+            @sio.on('connect')
+            def on_connect(sid, environ):
+                username = authenticate_user(environ)
+                if not username:
+                    return False
+                with sio.session(sid) as session:
+                    session['username'] = username
+
+            @sio.on('message')
+            def on_message(sid, msg):
+                with sio.session(sid) as session:
+                    print('received message from ', session['username'])
+        """
+        class _session_context_manager(object):
+            def __init__(self, server, sid, namespace):
+                self.server = server
+                self.sid = sid
+                self.namespace = namespace
+                self.session = None
+
+            def __enter__(self):
+                self.session = self.server.get_session(sid,
+                                                       namespace=namespace)
+                return self.session
+
+            def __exit__(self, *args):
+                self.server.save_session(sid, self.session,
+                                         namespace=namespace)
+
+        return _session_context_manager(self, sid, namespace)
+
     def disconnect(self, sid, namespace=None):
         """Disconnect a client.
 

tests/test_asyncio_namespace.py

@@ -170,6 +170,27 @@ def test_rooms(self):
         ns.rooms('sid', namespace='/bar')
         ns.server.rooms.assert_called_with('sid', namespace='/bar')
 
+    def test_session(self):
+        ns = asyncio_namespace.AsyncNamespace('/foo')
+        mock_server = mock.MagicMock()
+        mock_server.get_session = AsyncMock()
+        mock_server.save_session = AsyncMock()
+        ns._set_server(mock_server)
+        _run(ns.get_session('sid'))
+        ns.server.get_session.mock.assert_called_with('sid', namespace='/foo')
+        _run(ns.get_session('sid', namespace='/bar'))
+        ns.server.get_session.mock.assert_called_with('sid', namespace='/bar')
+        _run(ns.save_session('sid', {'a': 'b'}))
+        ns.server.save_session.mock.assert_called_with('sid', {'a': 'b'},
+                                                       namespace='/foo')
+        _run(ns.save_session('sid', {'a': 'b'}, namespace='/bar'))
+        ns.server.save_session.mock.assert_called_with('sid', {'a': 'b'},
+                                                       namespace='/bar')
+        ns.session('sid')
+        ns.server.session.assert_called_with('sid', namespace='/foo')
+        ns.session('sid', namespace='/bar')
+        ns.server.session.assert_called_with('sid', namespace='/bar')
+
     def test_disconnect(self):
         ns = asyncio_namespace.AsyncNamespace('/foo')
         mock_server = mock.MagicMock()

tests/test_asyncio_server.py

@@ -454,6 +454,38 @@ def test_send_with_ack_namespace(self, eio):
         _run(s._handle_eio_message('123', '3/foo,1["foo",2]'))
         cb.assert_called_once_with('foo', 2)
 
+    def test_session(self, eio):
+        fake_session = {}
+
+        async def fake_get_session(sid):
+            return fake_session
+
+        async def fake_save_session(sid, session):
+            global fake_session
+            fake_session = session
+
+        eio.return_value.send = AsyncMock()
+        s = asyncio_server.AsyncServer()
+        s.eio.get_session = fake_get_session
+        s.eio.save_session = fake_save_session
+
+        async def _test():
+            await s._handle_eio_connect('123', 'environ')
+            await s.save_session('123', {'foo': 'bar'})
+            async with s.session('123') as session:
+                self.assertEqual(session, {'foo': 'bar'})
+                session['foo'] = 'baz'
+                session['bar'] = 'foo'
+            self.assertEqual(await s.get_session('123'), {'foo': 'baz', 'bar': 'foo'})
+            self.assertEqual(fake_session, {'/': {'foo': 'baz', 'bar': 'foo'}})
+            async with s.session('123', namespace='/ns') as session:
+                self.assertEqual(session, {})
+                session['a'] = 'b'
+            self.assertEqual(await s.get_session('123', namespace='/ns'), {'a': 'b'})
+            self.assertEqual(fake_session, {'/': {'foo': 'baz', 'bar': 'foo'},
+                                            '/ns': {'a': 'b'}})
+        _run(_test())
+
     def test_disconnect(self, eio):
         eio.return_value.send = AsyncMock()
         s = asyncio_server.AsyncServer()

tests/test_namespace.py

@@ -120,6 +120,24 @@ def test_rooms(self):
         ns.rooms('sid', namespace='/bar')
         ns.server.rooms.assert_called_with('sid', namespace='/bar')
 
+    def test_session(self):
+        ns = namespace.Namespace('/foo')
+        ns._set_server(mock.MagicMock())
+        ns.get_session('sid')
+        ns.server.get_session.assert_called_with('sid', namespace='/foo')
+        ns.get_session('sid', namespace='/bar')
+        ns.server.get_session.assert_called_with('sid', namespace='/bar')
+        ns.save_session('sid', {'a': 'b'})
+        ns.server.save_session.assert_called_with('sid', {'a': 'b'},
+                                                  namespace='/foo')
+        ns.save_session('sid', {'a': 'b'}, namespace='/bar')
+        ns.server.save_session.assert_called_with('sid', {'a': 'b'},
+                                                  namespace='/bar')
+        ns.session('sid')
+        ns.server.session.assert_called_with('sid', namespace='/foo')
+        ns.session('sid', namespace='/bar')
+        ns.server.session.assert_called_with('sid', namespace='/bar')
+
     def test_disconnect(self):
         ns = namespace.Namespace('/foo')
         ns._set_server(mock.MagicMock())