Fixed #22634 -- Made the database-backed session backends more extensible.

Introduced an AbstractBaseSession model and hooks providing the option
of overriding the model class used by the session store and the session
store class used by the model.

3 files changed, 83 insertions, 9 deletions

diff --git a/tests/runtests.py b/tests/runtests.py
index 7367cd5d66..97c09c7372 100755
--- a/tests/runtests.py
+++ b/tests/runtests.py
@@ -140,6 +140,7 @@ def setup(verbosity, test_labels):
         # us skip creating migrations for the test models.
         'auth': 'django.contrib.auth.tests.migrations',
         'contenttypes': 'contenttypes_tests.migrations',
+        'sessions': 'sessions_tests.migrations',
     }
     log_config = DEFAULT_LOGGING
     # Filter out non-error logging so we don't have to capture it in lots of
diff --git a/tests/sessions_tests/custom_db_backend.py b/tests/sessions_tests/custom_db_backend.py
new file mode 100644
index 0000000000..5d4857e5bb
--- /dev/null
+++ b/tests/sessions_tests/custom_db_backend.py
@@ -0,0 +1,43 @@
+"""
+This custom Session model adds an extra column to store an account ID. In
+real-world applications, it gives you the option of querying the database for
+all active sessions for a particular account.
+"""
+from django.contrib.sessions.backends.db import SessionStore as DBStore
+from django.contrib.sessions.base_session import AbstractBaseSession
+from django.db import models
+
+
+class CustomSession(AbstractBaseSession):
+    """
+    A session model with a column for an account ID.
+    """
+    account_id = models.IntegerField(null=True, db_index=True)
+
+    class Meta:
+        app_label = 'sessions'
+
+    @classmethod
+    def get_session_store_class(cls):
+        return SessionStore
+
+
+class SessionStore(DBStore):
+    """
+    A database session store, that handles updating the account ID column
+    inside the custom session model.
+    """
+    @classmethod
+    def get_model_class(cls):
+        return CustomSession
+
+    def create_model_instance(self, data):
+        obj = super(SessionStore, self).create_model_instance(data)
+
+        try:
+            account_id = int(data.get('_auth_user_id'))
+        except (ValueError, TypeError):
+            account_id = None
+        obj.account_id = account_id
+
+        return obj
diff --git a/tests/sessions_tests/tests.py b/tests/sessions_tests/tests.py
index 76e625aa76..1a50720ffa 100644
--- a/tests/sessions_tests/tests.py
+++ b/tests/sessions_tests/tests.py
@@ -34,6 +34,8 @@ from django.utils import six, timezone
 from django.utils.encoding import force_text
 from django.utils.six.moves import http_cookies
 
+from .custom_db_backend import SessionStore as CustomDatabaseSession
+
 
 class SessionTestsMixin(object):
     # This does not inherit from TestCase to avoid any tests being run with this
@@ -355,6 +357,11 @@ class SessionTestsMixin(object):
 class DatabaseSessionTests(SessionTestsMixin, TestCase):
 
     backend = DatabaseSession
+    session_engine = 'django.contrib.sessions.backends.db'
+
+    @property
+    def model(self):
+        return self.backend.get_model_class()
 
     def test_session_str(self):
         "Session repr should be the session key."
@@ -362,7 +369,7 @@ class DatabaseSessionTests(SessionTestsMixin, TestCase):
         self.session.save()
 
         session_key = self.session.session_key
-        s = Session.objects.get(session_key=session_key)
+        s = self.model.objects.get(session_key=session_key)
 
         self.assertEqual(force_text(s), session_key)
 
@@ -374,7 +381,7 @@ class DatabaseSessionTests(SessionTestsMixin, TestCase):
         self.session['x'] = 1
         self.session.save()
 
-        s = Session.objects.get(session_key=self.session.session_key)
+        s = self.model.objects.get(session_key=self.session.session_key)
 
         self.assertEqual(s.get_decoded(), {'x': 1})
 
@@ -386,19 +393,18 @@ class DatabaseSessionTests(SessionTestsMixin, TestCase):
         self.session['y'] = 1
         self.session.save()
 
-        s = Session.objects.get(session_key=self.session.session_key)
+        s = self.model.objects.get(session_key=self.session.session_key)
         # Change it
-        Session.objects.save(s.session_key, {'y': 2}, s.expire_date)
+        self.model.objects.save(s.session_key, {'y': 2}, s.expire_date)
         # Clear cache, so that it will be retrieved from DB
         del self.session._session_cache
         self.assertEqual(self.session['y'], 2)
 
-    @override_settings(SESSION_ENGINE="django.contrib.sessions.backends.db")
     def test_clearsessions_command(self):
         """
         Test clearsessions command for clearing expired sessions.
         """
-        self.assertEqual(0, Session.objects.count())
+        self.assertEqual(0, self.model.objects.count())
 
         # One object in the future
         self.session['foo'] = 'bar'
@@ -412,10 +418,11 @@ class DatabaseSessionTests(SessionTestsMixin, TestCase):
         other_session.save()
 
         # Two sessions are in the database before clearsessions...
-        self.assertEqual(2, Session.objects.count())
-        management.call_command('clearsessions')
+        self.assertEqual(2, self.model.objects.count())
+        with override_settings(SESSION_ENGINE=self.session_engine):
+            management.call_command('clearsessions')
         # ... and one is deleted.
-        self.assertEqual(1, Session.objects.count())
+        self.assertEqual(1, self.model.objects.count())
 
 
 @override_settings(USE_TZ=True)
@@ -423,6 +430,29 @@ class DatabaseSessionWithTimeZoneTests(DatabaseSessionTests):
     pass
 
 
+class CustomDatabaseSessionTests(DatabaseSessionTests):
+    backend = CustomDatabaseSession
+    session_engine = 'sessions_tests.custom_db_backend'
+
+    def test_extra_session_field(self):
+        # Set the account ID to be picked up by a custom session storage
+        # and saved to a custom session model database column.
+        self.session['_auth_user_id'] = 42
+        self.session.save()
+
+        # Make sure that the customized create_model_instance() was called.
+        s = self.model.objects.get(session_key=self.session.session_key)
+        self.assertEqual(s.account_id, 42)
+
+        # Make the session "anonymous".
+        self.session.pop('_auth_user_id')
+        self.session.save()
+
+        # Make sure that save() on an existing session did the right job.
+        s = self.model.objects.get(session_key=self.session.session_key)
+        self.assertEqual(s.account_id, None)
+
+
 class CacheDBSessionTests(SessionTestsMixin, TestCase):
 
     backend = CacheDBSession