From 3b0409e572e5eca8246fe6ffbf0f7bfcc855edd1 Mon Sep 17 00:00:00 2001
From: SKairinos <stefan.kairinos@ocado.com>
Date: Thu, 7 Nov 2024 11:38:17 +0000
Subject: [PATCH] fix: session def

---
 codeforlife/models/abstract_base_session.py | 32 ++++++++++++++++-----
 codeforlife/user/models/session.py          |  9 +-----
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/codeforlife/models/abstract_base_session.py b/codeforlife/models/abstract_base_session.py
index a752275..d45dc54 100644
--- a/codeforlife/models/abstract_base_session.py
+++ b/codeforlife/models/abstract_base_session.py
@@ -5,7 +5,6 @@
 
 import typing as t
 
-from django.contrib.auth import get_user_model
 from django.contrib.sessions.base_session import (
     AbstractBaseSession as _AbstractBaseSession,
 )
@@ -15,6 +14,7 @@
 
 from .abstract_base_user import AbstractBaseUser
 
+# pylint: disable=duplicate-code
 if t.TYPE_CHECKING:
     from django_stubs_ext.db.models import TypedModelMeta
 
@@ -22,6 +22,9 @@
 else:
     TypedModelMeta = object
 
+AnyAbstractBaseUser = t.TypeVar("AnyAbstractBaseUser", bound=AbstractBaseUser)
+# pylint: enable=duplicate-code
+
 
 class AbstractBaseSession(_AbstractBaseSession):
     """
@@ -32,12 +35,6 @@ class AbstractBaseSession(_AbstractBaseSession):
     pk: str  # type: ignore[assignment]
 
     user_id: int
-    user = models.OneToOneField(
-        t.cast(t.Type[AbstractBaseUser], get_user_model()),
-        null=True,
-        blank=True,
-        on_delete=models.CASCADE,
-    )
 
     # pylint: disable-next=missing-class-docstring,too-few-public-methods
     class Meta(TypedModelMeta):
@@ -58,3 +55,24 @@ def store(self):
     @classmethod
     def get_session_store_class(cls) -> t.Type["BaseSessionStore"]:
         raise NotImplementedError
+
+    @staticmethod
+    def init_user_field(user_class: t.Type[AnyAbstractBaseUser]):
+        """Initializes the user field that relates a session to a user.
+
+        Example:
+            class Session(AbstractBaseSession):
+                user = AbstractBaseSession.init_user_field(User)
+
+        Args:
+            user_class: The user model to associate sessions to.
+
+        Returns:
+            A one-to-one field that relates to the provided user model.
+        """
+        return models.OneToOneField(
+            user_class,
+            null=True,
+            blank=True,
+            on_delete=models.CASCADE,
+        )
diff --git a/codeforlife/user/models/session.py b/codeforlife/user/models/session.py
index a6df0f9..a2819de 100644
--- a/codeforlife/user/models/session.py
+++ b/codeforlife/user/models/session.py
@@ -5,7 +5,6 @@
 
 import typing as t
 
-from django.db import models
 from django.db.models.query import QuerySet
 
 from ...models import AbstractBaseSession, BaseSessionStore
@@ -23,13 +22,7 @@ class Session(AbstractBaseSession):
 
     auth_factors: QuerySet["SessionAuthFactor"]
 
-    # TODO: remove in new schema
-    user = models.OneToOneField(  # type: ignore[assignment]
-        User,
-        null=True,
-        blank=True,
-        on_delete=models.CASCADE,
-    )
+    user = AbstractBaseSession.init_user_field(User)
 
     @classmethod
     def get_session_store_class(cls):