diff --git a/osf/management/commands/migrate_notifications.py b/osf/management/commands/migrate_notifications.py index c80565c8036..6e1786cd880 100644 --- a/osf/management/commands/migrate_notifications.py +++ b/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,38 +59,16 @@ 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, @@ -95,113 +76,102 @@ def timeout_handler(signum, frame): 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')