Rework notification subscriptions de-duplication command

Author: cslzchen

osf/management/commands/remove_duplicate_notification_subscriptions.py

@@ -7,8 +7,9 @@
 
 class Command(BaseCommand):
     help = (
-        'Remove duplicate NotificationSubscription records, keeping only '
-        'the highest-id record per (user, content_type, object_id, notification_type).'
+        'Remove duplicate NotificationSubscription records, keeping only the highest-id record: '
+        'Default uniqueness: (user, content_type, object_id, notification_type, is_digest); '
+        'Optional uniqueness with --exclude-is-digest: (user, content_type, object_id, notification_type).'
     )
 
     def add_arguments(self, parser):
@@ -17,39 +18,53 @@ def add_arguments(self, parser):
             action='store_true',
             help='Show how many rows would be deleted without deleting anything.',
         )
+        parser.add_argument(
+            '--exclude-is-digest',
+            action='store_true',
+            default=False,
+            help='Whether to exclude _is_digest field in unique_together')
 
     def handle(self, *args, **options):
+
         self.stdout.write('Finding duplicate NotificationSubscription records…')
 
-        to_remove = NotificationSubscription.objects.filter(
-            Exists(
-                NotificationSubscription.objects.filter(
-                    user_id=OuterRef('user_id'),
-                    content_type_id=OuterRef('content_type_id'),
-                    object_id=OuterRef('object_id'),
-                    notification_type_id=OuterRef('notification_type_id'),
-                    _is_digest=OuterRef('_is_digest'),
-                    id__gt=OuterRef('id'),  # keep most recent record
+        if options['exclude_is_digest']:
+            to_remove = NotificationSubscription.objects.filter(
+                Exists(
+                    NotificationSubscription.objects.filter(
+                        user_id=OuterRef('user_id'),
+                        content_type_id=OuterRef('content_type_id'),
+                        object_id=OuterRef('object_id'),
+                        notification_type_id=OuterRef('notification_type_id'),
+                        id__gt=OuterRef('id'),  # keep most recent record
+                    )
+                )
+            )
+        else:
+            to_remove = NotificationSubscription.objects.filter(
+                Exists(
+                    NotificationSubscription.objects.filter(
+                        user_id=OuterRef('user_id'),
+                        content_type_id=OuterRef('content_type_id'),
+                        object_id=OuterRef('object_id'),
+                        notification_type_id=OuterRef('notification_type_id'),
+                        _is_digest=OuterRef('_is_digest'),
+                        id__gt=OuterRef('id'),  # keep most recent record
+                    )
                 )
             )
-        )
 
         count = to_remove.count()
         self.stdout.write(f"Duplicates to remove: {count}")
 
-        if options['dry']:
-            self.stdout.write(
-                self.style.WARNING('Dry run enabled — no records were deleted.')
-            )
-            return
-
         if count == 0:
             self.stdout.write(self.style.SUCCESS('No duplicates found.'))
-            return
 
-        with transaction.atomic():
-            deleted, _ = to_remove.delete()
+        if options['dry']:
+            self.stdout.write(self.style.WARNING('Dry run enabled — no records were deleted.'))
+            return
 
-        self.stdout.write(
-            self.style.SUCCESS(f"Successfully removed {deleted} duplicate records.")
-        )
+        if count > 0:
+            with transaction.atomic():
+                deleted, _ = to_remove.delete()
+            self.stdout.write(self.style.SUCCESS(f"Successfully removed {deleted} duplicate records."))