From cd06088c77acf7212695d576f732c557fdea479b Mon Sep 17 00:00:00 2001
From: Alie Langston <alangsto@wellesley.edu>
Date: Thu, 31 Oct 2024 10:32:42 -0400
Subject: [PATCH] feat: add management command to remove expired messages

---
 CHANGELOG.rst                                 |  1 +
 .../commands/retire_user_messages.py          | 68 ++++++++++++++++++
 .../tests/test_retire_user_messages.py        | 69 +++++++++++++++++++
 tests/__init__.py                             |  0
 4 files changed, 138 insertions(+)
 create mode 100644 learning_assistant/management/commands/retire_user_messages.py
 create mode 100644 learning_assistant/management/commands/tests/test_retire_user_messages.py
 create mode 100644 tests/__init__.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index fb53e08..8a3b82f 100644
--- a/CHANGELOG.rst
+++ b/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
 ******************
diff --git a/learning_assistant/management/commands/retire_user_messages.py b/learning_assistant/management/commands/retire_user_messages.py
new file mode 100644
index 0000000..12862cf
--- /dev/null
+++ b/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.')
diff --git a/learning_assistant/management/commands/tests/test_retire_user_messages.py b/learning_assistant/management/commands/tests/test_retire_user_messages.py
new file mode 100644
index 0000000..6d7062a
--- /dev/null
+++ b/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='tester@test.com')
+        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)
+
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29