feat: add management command to remove expired messages

CHANGELOG.rst

@@ -14,6 +14,7 @@ Change Log
 Unreleased
 **********
 * Add LearningAssistantMessage model
+* Add management command to remove expired messages
 
 4.3.3 - 2024-10-15
 ******************

learning_assistant/management/commands/retire_user_messages.py

@@ -0,0 +1,68 @@
+""""
+Django management command to remove LearningAssistantMessage objects
+if they have reached their expiration date.
+"""
+from datetime import datetime, timedelta
+import logging
+import time
+
+from django.conf import settings
+from django.core.management.base import BaseCommand
+
+from learning_assistant.models import LearningAssistantMessage
+
+log = logging.getLogger(__name__)
+
+
+class Command(BaseCommand):
+    """
+    Django Management command to remove expired messages.
+    """
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--batch_size',
+            action='store',
+            dest='batch_size',
+            type=int,
+            default=300,
+            help='Maximum number of messages to remove. '
+                 'This helps avoid overloading the database while updating large amount of data.'
+        )
+        parser.add_argument(
+            '--sleep_time',
+            action='store',
+            dest='sleep_time',
+            type=int,
+            default=10,
+            help='Sleep time in seconds between update of batches'
+        )
+
+    def handle(self, *args, **options):
+        """
+        Management command entry point.
+        """
+        batch_size = options['batch_size']
+        sleep_time = options['sleep_time']
+
+        expiry_date = datetime.now() - timedelta(days=getattr(settings, 'LEARNING_ASSISTANT_MESSAGES_EXPIRY', 30))
+
+        total_deleted = 0
+        deleted_count = None
+
+        while deleted_count != 0:
+            ids_to_delete = LearningAssistantMessage.objects.filter(
+                created__lte=expiry_date
+            ).values_list('id', flat=True)[:batch_size]
+
+            ids_to_delete = list(ids_to_delete)
+            delete_queryset = LearningAssistantMessage.objects.filter(
+                id__in=ids_to_delete
+            )
+            deleted_count, _ = delete_queryset.delete()
+
+            total_deleted += deleted_count
+            log.info(f'{deleted_count} messages deleted.')
+            time.sleep(sleep_time)
+
+        log.info(f'Job completed. {total_deleted} messages deleted.')

learning_assistant/management/commands/tests/test_retire_user_messages.py

@@ -0,0 +1,69 @@
+"""
+Tests for the retire_user_messages management command
+"""
+from datetime import datetime, timedelta
+
+from django.contrib.auth import get_user_model
+from django.core.management import call_command
+from django.test import TestCase
+
+from learning_assistant.models import LearningAssistantMessage
+
+User = get_user_model()
+
+
+class RetireUserMessagesTests(TestCase):
+    """
+    Tests for the retire_user_messages command.
+    """
+
+    def setUp(self):
+        """
+        Build up test data
+        """
+        super().setUp()
+        self.user = User(username='tester', email='[email protected]')
+        self.user.save()
+
+        self.course_id = 'course-v1:edx+test+23'
+
+        LearningAssistantMessage.objects.create(
+            user=self.user,
+            course_id=self.course_id,
+            role='user',
+            content='Hello',
+            created=datetime.now() - timedelta(days=60)
+        )
+
+        LearningAssistantMessage.objects.create(
+            user=self.user,
+            course_id=self.course_id,
+            role='user',
+            content='Hello',
+            created=datetime.now() - timedelta(days=2)
+        )
+
+        LearningAssistantMessage.objects.create(
+            user=self.user,
+            course_id=self.course_id,
+            role='user',
+            content='Hello',
+            created=datetime.now() - timedelta(days=4)
+        )
+
+    def test_run_command(self):
+        """
+        Run the management command
+        """
+        current_messages = LearningAssistantMessage.objects.filter()
+        self.assertEqual(len(current_messages), 3)
+
+        call_command(
+            'retire_user_messages',
+            batch_size=2,
+            sleep_time=0,
+        )
+
+        current_messages = LearningAssistantMessage.objects.filter()
+        self.assertEqual(len(current_messages), 2)
+