[ENG-8823] Migration Improvements for notification refator

osf/management/commands/migrate_notifications.py

@@ -2,20 +2,27 @@
 import time
 import signal
 from contextlib import contextmanager
+
 from django.contrib.contenttypes.models import ContentType
 from django.core.management.base import BaseCommand, CommandError
 from django.db import transaction
+
 from osf.models import NotificationType, NotificationSubscription
 from osf.models.notifications import NotificationSubscriptionLegacy
 from osf.management.commands.populate_notification_types import populate_notification_types
+from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
+TIMEOUT_SECONDS = 3600  # 60 minutes timeout
+BATCH_SIZE = 1000       # Default batch size
+
 FREQ_MAP = {
     'none': 'none',
     'email_digest': 'weekly',
     'email_transactional': 'instantly',
 }
+
 EVENT_NAME_TO_NOTIFICATION_TYPE = {
     # Provider notifications
     'new_pending_withdraw_requests': NotificationType.Type.PROVIDER_NEW_PENDING_WITHDRAW_REQUESTS,
@@ -40,10 +47,6 @@
 }
 
 
-TIMEOUT_SECONDS = 3600  # 60 minutes timeout
-BATCH_SIZE = 1000  # batch size, can be changed
-
-
 @contextmanager
 def time_limit(seconds):
     def signal_handler(signum, frame):
@@ -56,152 +59,119 @@ def signal_handler(signum, frame):
     finally:
         signal.alarm(0)
 
-def migrate_legacy_notification_subscriptions(
-    dry_run=False,
-    batch_size=1000,
-    timeout_per_batch=300,
-    default_frequency='none'
-):
 
-    logger.info('Starting legacy notification subscription migration...')
+def iter_batches(first_id: int, last_id: int, batch_size: int):
+    """Yield [start_id, end_id] ranges for batching."""
+    for start in range(first_id, last_id + 1, batch_size):
+        yield start, min(start + batch_size - 1, last_id)
 
-    PROVIDER_BASED_LEGACY_NOTIFICATION_TYPES = [
-        'new_pending_submissions',
-        'new_pending_withdraw_requests',
-        'reviews_submission_confirmation',
-        'reviews_moderator_submission_confirmation',
-        'reviews_reject_confirmation',
-        'reviews_accept_confirmation',
-        'reviews_resubmission_confirmation',
-        'reviews_comment_edited',
-        'contributor_added_preprint',
-        'confirm_email_moderation',
-        'moderator_added',
-        'confirm_email_preprints',
-        'user_invite_preprint',
-    ]
-    def timeout_handler(signum, frame):
-        raise TimeoutError('Batch processing timed out')
-
-    # Notification type IDs
-    notiftype_map = dict(NotificationType.objects.values_list('name', 'id'))
 
-    # Cache existing keys
-    existing_keys = set(
+def build_existing_keys():
+    """Fetch already migrated subscription keys to prevent duplicates."""
+    return set(
         (
             user_id,
             content_type_id,
             int(object_id) if object_id else None,
             notification_type_id,
         )
         for user_id, content_type_id, object_id, notification_type_id in
-        NotificationSubscription.objects.all().values_list(
+        NotificationSubscription.objects.values_list(
             'user_id', 'content_type_id', 'object_id', 'notification_type_id'
         )
     )
 
-    total = NotificationSubscriptionLegacy.objects.count()
-    created = 0
-    skipped = 0
-    last_id = 0
+
+def migrate_legacy_notification_subscriptions(
+    dry_run=False,
+    batch_size=BATCH_SIZE,
+    default_frequency='none',
+    start_id=0,
+):
+    logger.info('Starting legacy notification subscription migration...')
+
+    legacy_qs = NotificationSubscriptionLegacy.objects.filter(id__gte=start_id).order_by('id')
+    total = legacy_qs.count()
+    if total == 0:
+        logger.info('No legacy subscriptions to migrate.')
+        return
+
+    notiftype_map = dict(NotificationType.objects.values_list('name', 'id'))
+    existing_keys = build_existing_keys()
+
+    created, skipped = 0, 0
+    content_type_cache = {}
+
+    first_id, last_id = legacy_qs.first().id, legacy_qs.last().id
     start_time_total = time.time()
 
-    while True:
-        batch_start_time = time.time()
-        # Fetch the next chunk directly from DB
+    for batch_range in tqdm(list(iter_batches(first_id, last_id, batch_size)), desc='Processing', unit='batch'):
         batch = list(
             NotificationSubscriptionLegacy.objects
-            .filter(id__gt=last_id)
-            .order_by('id')[:batch_size]
+            .filter(id__range=batch_range)
+            .order_by('id')
+            .select_related('provider', 'node', 'user')
         )
         if not batch:
-            break
-        subscriptions_to_create = []
+            continue
 
-        signal.signal(signal.SIGALRM, timeout_handler)
-        signal.alarm(timeout_per_batch)
+        subscriptions_to_create = []
 
-        try:
-            for legacy in batch:
-                event_name = legacy.event_name
-                if event_name in PROVIDER_BASED_LEGACY_NOTIFICATION_TYPES:
-                    subscribed_object = legacy.provider
-                    if not subscribed_object:
-                        skipped += 1
-                        continue
-                elif legacy.node:
-                    subscribed_object = legacy.node
-                elif legacy.user:
-                    subscribed_object = legacy.user
-                else:
-                    skipped += 1
-                    continue
-
-                content_type = ContentType.objects.get_for_model(subscribed_object.__class__)
-                notif_enum = EVENT_NAME_TO_NOTIFICATION_TYPE.get(event_name)
-                if not notif_enum:
-                    skipped += 1
-                    continue
-
-                notification_type_id = notiftype_map.get(notif_enum)
-                if not notification_type_id:
-                    skipped += 1
-                    continue
-
-                key = (
-                    legacy.user_id,
-                    content_type.id,
-                    int(subscribed_object.id),
-                    notification_type_id,
-                )
-                if key in existing_keys:
-                    skipped += 1
-                    continue
-
-                frequency = 'weekly' if getattr(legacy, 'email_digest', False) else default_frequency
-
-                if dry_run:
-                    created += 1
-                else:
-                    subscriptions_to_create.append(NotificationSubscription(
-                        notification_type_id=notification_type_id,
-                        user_id=legacy.user_id,
-                        content_type=content_type,
-                        object_id=subscribed_object.id,
-                        message_frequency=frequency,
-                    ))
-                    existing_keys.add(key)
-
-            if not dry_run and subscriptions_to_create:
-                with transaction.atomic():
-                    NotificationSubscription.objects.bulk_create(
-                        subscriptions_to_create,
-                        ignore_conflicts=True,
-                    )
-                    created += len(subscriptions_to_create)
-
-            # Logging ETA
-            batch_time = time.time() - batch_start_time
-            processed = last_id + len(batch)
-            rate = processed / (time.time() - start_time_total)
-            eta = (total - processed) / rate if rate else 0
-
-            logger.info(
-                f"Processed batch {last_id}-{last_id + len(batch)} "
-                f"in {batch_time:.2f}s | "
-                f"Progress {processed}/{total} ({processed / total:.1%}) | "
-                f"ETA ~ {eta / 60:.1f} min"
+        for legacy in batch:
+            event_name = legacy.event_name
+            subscribed_object = legacy.provider or legacy.node or legacy.user
+            if not subscribed_object:
+                skipped += 1
+                continue
+
+            model_class = subscribed_object.__class__
+            if model_class not in content_type_cache:
+                content_type_cache[model_class] = ContentType.objects.get_for_model(model_class)
+            content_type = content_type_cache[model_class]
+
+            notif_enum = EVENT_NAME_TO_NOTIFICATION_TYPE.get(event_name)
+            if not notif_enum:
+                skipped += 1
+                continue
+
+            notification_type_id = notiftype_map.get(notif_enum)
+            if not notification_type_id:
+                skipped += 1
+                continue
+
+            key = (
+                legacy.user_id,
+                content_type.id,
+                int(subscribed_object.id),
+                notification_type_id,
             )
+            if key in existing_keys:
+                skipped += 1
+                continue
+
+            frequency = 'weekly' if getattr(legacy, 'email_digest', False) else default_frequency
+
+            if dry_run:
+                created += 1
+            else:
+                subscriptions_to_create.append(NotificationSubscription(
+                    notification_type_id=notification_type_id,
+                    user_id=legacy.user_id,
+                    content_type=content_type,
+                    object_id=subscribed_object.id,
+                    message_frequency=frequency,
+                ))
+                existing_keys.add(key)
+
+        if not dry_run and subscriptions_to_create:
+            with transaction.atomic():
+                NotificationSubscription.objects.bulk_create(
+                    subscriptions_to_create,
+                    ignore_conflicts=True,
+                )
+            created += len(subscriptions_to_create)
 
-        except TimeoutError:
-            logger.error(f"Batch {last_id}-{last_id + len(batch)} timed out, skipping.")
-            skipped += len(batch)
-        except Exception as e:
-            logger.exception(f"Batch {last_id}-{last_id + len(batch)} failed: {e}")
-            skipped += len(batch)
-        finally:
-            signal.alarm(0)
-            last_id = batch[-1].id
+        logger.info(f"Processed batch {batch_range[0]}-{batch_range[1]} (Created: {created}, Skipped: {skipped})")
 
     elapsed_total = time.time() - start_time_total
     logger.info(
@@ -210,6 +180,16 @@ def timeout_handler(signum, frame):
     )
 
 
+def run_migration(dry_run: bool, batch_size: int, start_id: int):
+    """Main entry point for command and tests."""
+    with time_limit(TIMEOUT_SECONDS):
+        if not dry_run:
+            with transaction.atomic():
+                logger.info('Populating notification types...')
+                populate_notification_types(None, {})
+        migrate_legacy_notification_subscriptions(dry_run=dry_run, batch_size=batch_size, start_id=start_id)
+
+
 class Command(BaseCommand):
     help = 'Migrate legacy NotificationSubscriptionLegacy objects to new Notification app models.'
 
@@ -220,19 +200,27 @@ def add_arguments(self, parser):
             help='Run migration in dry-run mode (no DB changes will be committed).'
         )
 
-    def handle(self, *args, **options):
-        dry_run = options['dry_run']
-
-        try:
-            with time_limit(TIMEOUT_SECONDS):
-                if not dry_run:
-                    with transaction.atomic():
-                        logger.info('Populating notification types...')
-                        populate_notification_types(args, options)
+        parser.add_argument(
+            '--batch-size',
+            type=int,
+            default=BATCH_SIZE,
+            help=f'Batch size (default: {BATCH_SIZE})'
+        )
 
-                with transaction.atomic():
-                    migrate_legacy_notification_subscriptions(dry_run=dry_run)
+        parser.add_argument(
+            '--start-id',
+            type=int,
+            default=0,
+            help='Start migrating from this ID'
+        )
 
+    def handle(self, *args, **options):
+        try:
+            run_migration(
+                dry_run=options['dry_run'],
+                batch_size=options['batch_size'],
+                start_id=options['start_id'],
+            )
         except TimeoutError:
             logger.error('Migration timed out. Rolling back changes.')
             raise CommandError('Migration failed due to timeout')