Robust auto-cleaning functionality

src/reactpy_django/clean.py

@@ -0,0 +1,88 @@
+import logging
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING
+
+from django.contrib.auth import get_user_model
+from django.utils import timezone
+
+from reactpy_django.utils import get_pk
+
+_logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from reactpy_django.models import Config
+
+
+def clean_all(immediate: bool = False, manual_clean=True):
+    from reactpy_django.config import (
+        REACTPY_AUTO_CLEAN_SESSIONS,
+        REACTPY_AUTO_CLEAN_USER_DATA,
+    )
+    from reactpy_django.models import Config
+
+    config = Config.load()
+    if immediate or is_clean_needed(config):
+        if manual_clean or REACTPY_AUTO_CLEAN_SESSIONS:
+            clean_sessions()
+        if manual_clean or REACTPY_AUTO_CLEAN_USER_DATA:
+            clean_user_data()
+        config.cleaned_at = timezone.now()
+        config.save()
+
+
+def clean_sessions():
+    """Deletes expired component sessions from the database.
+    As a performance optimization, this is only run once every REACTPY_SESSION_MAX_AGE seconds.
+    """
+    from reactpy_django.config import REACTPY_DEBUG_MODE, REACTPY_SESSION_MAX_AGE
+    from reactpy_django.models import ComponentSession
+
+    start_time = timezone.now()
+    expiration_date = timezone.now() - timedelta(seconds=REACTPY_SESSION_MAX_AGE)
+    ComponentSession.objects.filter(last_accessed__lte=expiration_date).delete()
+
+    if REACTPY_DEBUG_MODE:
+        inspect_clean_duration(start_time, "component sessions")
+
+
+def clean_user_data():
+    """Delete any user data that is not associated with an existing `User`.
+    This is a safety measure to ensure that we don't have any orphaned data in the database.
+
+    Our `UserDataModel` is supposed to automatically get deleted on every `User` delete signal.
+    However, we can't use Django to enforce this relationship since ReactPy can be configured to
+    use any database.
+    """
+    from reactpy_django.config import REACTPY_DEBUG_MODE
+    from reactpy_django.models import UserDataModel
+
+    start_time = timezone.now()
+    user_model = get_user_model()
+    all_users = user_model.objects.all()
+    all_user_pks = all_users.values_list(user_model._meta.pk.name, flat=True)  # type: ignore
+    UserDataModel.objects.exclude(user_pk__in=all_user_pks).delete()
+
+    if REACTPY_DEBUG_MODE:
+        inspect_clean_duration(start_time, "user data")
+
+
+def is_clean_needed(config: Config | None = None) -> bool:
+    from reactpy_django.config import REACTPY_SESSION_MAX_AGE
+    from reactpy_django.models import Config
+
+    config = config or Config.load()
+    cleaned_at = config.cleaned_at
+    clean_needed_by = cleaned_at + timedelta(seconds=REACTPY_SESSION_MAX_AGE)
+
+    return timezone.now() >= clean_needed_by
+
+
+def inspect_clean_duration(start_time: datetime, task_name: str):
+    clean_duration = timezone.now() - start_time
+    if clean_duration.total_seconds() > 1:
+        _logger.warning(
+            "ReactPy has taken %s seconds to clean %s. "
+            "This may indicate a performance issue with your system, cache, or database.",
+            clean_duration.total_seconds(),
+            task_name,
+        )

src/reactpy_django/config.py

@@ -118,3 +118,18 @@
     "REACTPY_AUTO_RELOGIN",
     False,
 )
+REACTPY_AUTO_CLEAN_INTERVAL: int = getattr(
+    settings,
+    "REACTPY_AUTO_CLEAN_INTERVAL",
+    604800,  # Default to 7 days
+)
+REACTPY_AUTO_CLEAN_SESSIONS: bool = getattr(
+    settings,
+    "REACTPY_AUTO_CLEAN_SESSIONS",
+    True,
+)
+REACTPY_AUTO_CLEAN_USER_DATA: bool = getattr(
+    settings,
+    "REACTPY_AUTO_CLEAN_USER_DATA",
+    True,
+)

src/reactpy_django/hooks.py

@@ -33,7 +33,7 @@
     QueryOptions,
     UserData,
 )
-from reactpy_django.utils import generate_obj_name, get_user_pk
+from reactpy_django.utils import generate_obj_name, get_pk
 
 if TYPE_CHECKING:
     from django.contrib.auth.models import AbstractUser
@@ -344,7 +344,7 @@ async def _set_user_data(data: dict):
         if user.is_anonymous:
             raise ValueError("AnonymousUser cannot have user data.")
 
-        pk = get_user_pk(user)
+        pk = get_pk(user)
         model, _ = await UserDataModel.objects.aget_or_create(user_pk=pk)
         model.data = pickle.dumps(data)
         await model.asave()
@@ -386,7 +386,7 @@ async def _get_user_data(
     if not user or user.is_anonymous:
         return None
 
-    pk = get_user_pk(user)
+    pk = get_pk(user)
     model, _ = await UserDataModel.objects.aget_or_create(user_pk=pk)
     data = pickle.loads(model.data) if model.data else {}
 

src/reactpy_django/utils.py

@@ -6,20 +6,17 @@
 import os
 import re
 from asyncio import iscoroutinefunction
-from datetime import timedelta
 from fnmatch import fnmatch
 from importlib import import_module
 from typing import Any, Callable, Sequence
 
-import orjson as pickle
 from asgiref.sync import async_to_sync
 from channels.db import database_sync_to_async
 from django.db.models import ManyToManyField, ManyToOneRel, prefetch_related_objects
 from django.db.models.base import Model
 from django.db.models.query import QuerySet
 from django.http import HttpRequest, HttpResponse
 from django.template import engines
-from django.utils import timezone
 from django.utils.encoding import smart_str
 from django.views import View
 from reactpy.core.layout import Layout
@@ -351,36 +348,6 @@ def create_cache_key(*args):
     return f"reactpy_django:{':'.join(str(arg) for arg in args)}"
 
 
-def delete_expired_sessions(immediate: bool = False):
-    """Deletes expired component sessions from the database.
-    As a performance optimization, this is only run once every REACTPY_SESSION_MAX_AGE seconds.
-    """
-    from .config import REACTPY_DEBUG_MODE, REACTPY_SESSION_MAX_AGE
-    from .models import ComponentSession, Config
-
-    config = Config.load()
-    start_time = timezone.now()
-    cleaned_at = config.cleaned_at
-    clean_needed_by = cleaned_at + timedelta(seconds=REACTPY_SESSION_MAX_AGE)
-
-    # Delete expired component parameters
-    if immediate or timezone.now() >= clean_needed_by:
-        expiration_date = timezone.now() - timedelta(seconds=REACTPY_SESSION_MAX_AGE)
-        ComponentSession.objects.filter(last_accessed__lte=expiration_date).delete()
-        config.cleaned_at = timezone.now()
-        config.save()
-
-    # Check if cleaning took abnormally long
-    if REACTPY_DEBUG_MODE:
-        clean_duration = timezone.now() - start_time
-        if clean_duration.total_seconds() > 1:
-            _logger.warning(
-                "ReactPy has taken %s seconds to clean up expired component sessions. "
-                "This may indicate a performance issue with your system, cache, or database.",
-                clean_duration.total_seconds(),
-            )
-
-
 class SyncLayout(Layout):
     """Sync adapter for ReactPy's `Layout`. Allows it to be used in Django template tags.
     This can be removed when Django supports async template tags.
@@ -397,7 +364,6 @@ def render(self):
         return async_to_sync(super().render)()
 
 
-def get_user_pk(user, serialize=False):
-    """Returns the primary key value for a user model instance."""
-    pk = getattr(user, user._meta.pk.name)
-    return pickle.dumps(pk) if serialize else pk
+def get_pk(model):
+    """Returns the value of the primary key for a Django model."""
+    return getattr(model, model._meta.pk.name)

src/reactpy_django/websocket/consumer.py

@@ -23,8 +23,8 @@
 from reactpy.core.layout import Layout
 from reactpy.core.serve import serve_layout
 
+from reactpy_django.clean import clean_all
 from reactpy_django.types import ComponentParams
-from reactpy_django.utils import delete_expired_sessions
 
 if TYPE_CHECKING:
     from django.contrib.auth.models import AbstractUser
@@ -96,27 +96,29 @@ async def connect(self) -> None:
 
     async def disconnect(self, code: int) -> None:
         """The browser has disconnected."""
+        from reactpy_django.config import REACTPY_AUTO_CLEAN_INTERVAL
+
         self.dispatcher.cancel()
 
+        # Update the component's last_accessed timestamp
         if self.component_session:
-            # Clean up expired component sessions
             try:
-                await database_sync_to_async(delete_expired_sessions)()
+                await self.component_session.asave()
             except Exception:
                 await asyncio.to_thread(
                     _logger.error,
-                    "ReactPy has failed to delete expired component sessions!\n"
+                    "ReactPy has failed to save component session!\n"
                     f"{traceback.format_exc()}",
                 )
 
-            # Update the last_accessed timestamp
+        # Queue a cleanup, if needed
+        if REACTPY_AUTO_CLEAN_INTERVAL > -1:
             try:
-                await self.component_session.asave()
+                await database_sync_to_async(clean_all)()
             except Exception:
                 await asyncio.to_thread(
                     _logger.error,
-                    "ReactPy has failed to save component session!\n"
-                    f"{traceback.format_exc()}",
+                    "ReactPy cleaning failed!\n" f"{traceback.format_exc()}",
                 )
 
         await super().disconnect(code)

tests/test_app/tests/test_database.py

@@ -4,8 +4,8 @@
 
 import dill as pickle
 from django.test import TransactionTestCase
-from reactpy_django import utils
-from reactpy_django.models import ComponentSession
+from reactpy_django import clean
+from reactpy_django.models import ComponentSession, UserDataModel
 from reactpy_django.types import ComponentParams
 
 
@@ -17,7 +17,7 @@ class RoutedDatabaseTests(TransactionTestCase):
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
-        utils.delete_expired_sessions(immediate=True)
+        clean.clean_sessions(immediate=True)
 
     def test_component_params(self):
         # Make sure the ComponentParams table is empty
@@ -27,7 +27,7 @@ def test_component_params(self):
         # Check if a component params are in the database
         self.assertEqual(ComponentSession.objects.count(), 1)
         self.assertEqual(
-            pickle.loads(ComponentSession.objects.first().params), params_1  # type: ignore
+            pickle.loads(ComponentSession.objects.first().params), params_1
         )
 
         # Force `params_1` to expire
@@ -41,14 +41,43 @@ def test_component_params(self):
         self.assertEqual(ComponentSession.objects.count(), 2)
 
         # Delete the first component params based on expiration time
-        utils.delete_expired_sessions()  # Don't use `immediate` to test timestamping logic
+        clean.clean_sessions()  # Don't use `immediate` to test timestamping logic
 
         # Make sure `params_1` has expired
         self.assertEqual(ComponentSession.objects.count(), 1)
         self.assertEqual(
-            pickle.loads(ComponentSession.objects.first().params), params_2  # type: ignore
+            pickle.loads(ComponentSession.objects.first().params), params_2
         )
 
+    def test_user_data_cleanup(self):
+        from django.contrib.auth.models import User
+
+        # Create UserData for real user #1
+        user = User.objects.create_user(username=uuid4().hex, password=uuid4().hex)
+        user_data = UserDataModel(user_pk=user.pk)
+        user_data.save()
+
+        # Create UserData for real user #2
+        user = User.objects.create_user(username=uuid4().hex, password=uuid4().hex)
+        user_data = UserDataModel(user_pk=user.pk)
+        user_data.save()
+
+        # Keep track of the count of UserData objects
+        original_count = UserDataModel.objects.count()
+
+        # Create UserData for a fake user (orphaned)
+        user_data = UserDataModel(user_pk=uuid4().hex)
+        user_data.save()
+
+        # Make sure the orphaned user data object is deleted
+        self.assertNotEqual(UserDataModel.objects.count(), original_count)
+        clean.clean_user_data()
+        self.assertEqual(UserDataModel.objects.count(), original_count)
+
+        # Check if deleting a user deletes the associated UserData
+        user.delete()
+        self.assertEqual(UserDataModel.objects.count(), original_count - 1)
+
     def _save_params_to_db(self, value: Any) -> ComponentParams:
         db = list(self.databases)[0]
         param_data = ComponentParams((value,), {"test_value": value})