Update order_by and add unit test for concurrency

Author: rajpatel24

kpi/tasks.py

@@ -94,7 +94,7 @@ def cleanup_anonymous_exports(**kwargs):
         date_created__lt=cutoff_time,
     ).exclude(
         status=ImportExportStatusChoices.PROCESSING
-    ).order_by('id').values_list('pk', flat=True)[:BATCH_SIZE]
+    ).order_by('date_created').values_list('pk', flat=True)[:BATCH_SIZE]
 
     if not old_export_ids:
         logging.info('No old anonymous exports to clean up.')

kpi/tests/test_cleanup_anonymous_exports.py

@@ -1,7 +1,10 @@
+import threading
+import time
 from datetime import timedelta
 
+from django.db import transaction
 from django.utils import timezone
-from django.test import TestCase
+from django.test import TransactionTestCase
 
 from kpi.models.import_export_task import (
     ImportExportStatusChoices,
@@ -11,7 +14,7 @@
 from kpi.utils.object_permission import get_anonymous_user
 
 
-class AnonymousExportCleanupTestCase(TestCase):
+class AnonymousExportCleanupTestCase(TransactionTestCase):
     def _create_export_task(
         self, status=ImportExportStatusChoices.COMPLETE, minutes_old=60
     ):
@@ -59,3 +62,52 @@ def test_processing_exports_are_not_deleted(self):
                 uid=processing_export.uid
             ).exists()
         )
+
+    def test_concurrency_locked_rows_are_skipped(self):
+        exports = [self._create_export_task(minutes_old=120) for _ in range(5)]
+        export_pks = [e.pk for e in exports]
+        locked_pks = export_pks[:3]
+        unlocked_pks = export_pks[3:]
+
+        def lock_and_hold():
+            """
+            Acquire lock on first 3 exports and hold for 5 seconds,
+            this will block cleanup from acquiring the lock on these rows and
+            in the meantime other rows should be cleaned up
+            """
+            with transaction.atomic():
+                # Acquire lock on first 3 exports
+                list(
+                    SubmissionExportTask.objects
+                    .select_for_update()
+                    .filter(pk__in=locked_pks)
+                )
+
+                # Hold lock for 5 seconds,
+                # during this time cleanup should run
+                time.sleep(5)
+
+        # Start thread that will hold the lock
+        lock_thread = threading.Thread(target=lock_and_hold, daemon=True)
+        lock_thread.start()
+
+        # Give thread time to acquire lock
+        time.sleep(1)
+
+        # Run cleanup while lock is held
+        cleanup_anonymous_exports()
+
+        # Wait for thread to finish
+        lock_thread.join(timeout=10)
+
+        # Verify locked rows were not deleted
+        remaining_locked = SubmissionExportTask.objects.filter(
+            pk__in=locked_pks
+        ).count()
+        self.assertEqual(remaining_locked, 3)
+
+        # Verify unlocked rows were deleted
+        remaining_unlocked = SubmissionExportTask.objects.filter(
+            pk__in=unlocked_pks
+        ).count()
+        self.assertEqual(remaining_unlocked, 0)