Skip to content

Commit 873f28d

Browse files
bodintsovbodintsov
committed
refactored migrate_notifications.py
1 parent 8bb9287 commit 873f28d

2 files changed

Lines changed: 347 additions & 70 deletions

File tree

‎osf/management/commands/migrate_notifications.py‎

Lines changed: 213 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
import logging
2+
import time
3+
import signal
4+
from contextlib import contextmanager
25
from django.contrib.contenttypes.models import ContentType
6+
from django.core.management.base import BaseCommand, CommandError
7+
from django.db import transaction
38
from osf.models import NotificationType, NotificationSubscription
49
from osf.models.notifications import NotificationSubscriptionLegacy
5-
from django.core.management.base import BaseCommand
6-
from django.db import transaction
710
from osf.management.commands.populate_notification_types import populate_notification_types
811

912
logger = logging.getLogger(__name__)
@@ -14,49 +17,226 @@
1417
'email_transactional': 'instantly',
1518
}
1619
EVENT_NAME_TO_NOTIFICATION_TYPE = {
20+
# Provider notifications
21+
'new_pending_withdraw_requests': NotificationType.Type.PROVIDER_NEW_PENDING_WITHDRAW_REQUESTS,
22+
'contributor_added_preprint': NotificationType.Type.PROVIDER_CONTRIBUTOR_ADDED_PREPRINT,
1723
'new_pending_submissions': NotificationType.Type.PROVIDER_NEW_PENDING_SUBMISSIONS,
24+
'moderator_added': NotificationType.Type.PROVIDER_MODERATOR_ADDED,
25+
'reviews_submission_confirmation': NotificationType.Type.PROVIDER_REVIEWS_SUBMISSION_CONFIRMATION,
26+
'reviews_resubmission_confirmation': NotificationType.Type.PROVIDER_REVIEWS_RESUBMISSION_CONFIRMATION,
27+
'confirm_email_moderation': NotificationType.Type.PROVIDER_CONFIRM_EMAIL_MODERATION,
28+
29+
# Node notifications
1830
'file_updated': NotificationType.Type.NODE_FILE_UPDATED,
19-
'comments': None,
31+
32+
# Collection submissions
33+
'collection_submission_submitted': NotificationType.Type.COLLECTION_SUBMISSION_SUBMITTED,
34+
'collection_submission_accepted': NotificationType.Type.COLLECTION_SUBMISSION_ACCEPTED,
35+
'collection_submission_rejected': NotificationType.Type.COLLECTION_SUBMISSION_REJECTED,
36+
'collection_submission_removed_admin': NotificationType.Type.COLLECTION_SUBMISSION_REMOVED_ADMIN,
37+
'collection_submission_removed_moderator': NotificationType.Type.COLLECTION_SUBMISSION_REMOVED_MODERATOR,
38+
'collection_submission_removed_private': NotificationType.Type.COLLECTION_SUBMISSION_REMOVED_PRIVATE,
39+
'collection_submission_cancel': NotificationType.Type.COLLECTION_SUBMISSION_CANCEL,
2040
}
2141

22-
def migrate_legacy_notification_subscriptions(*args, **kwargs):
23-
"""
24-
Migrate legacy NotificationSubscription data to new notifications app.
25-
"""
26-
logger.info('Beginning legacy notification subscription migration...')
42+
43+
TIMEOUT_SECONDS = 3600 # 60 minutes timeout
44+
BATCH_SIZE = 1000
45+
46+
47+
@contextmanager
48+
def time_limit(seconds):
49+
def signal_handler(signum, frame):
50+
raise TimeoutError("Migration timed out")
51+
52+
signal.signal(signal.SIGALRM, signal_handler)
53+
signal.alarm(seconds)
54+
try:
55+
yield
56+
finally:
57+
signal.alarm(0)
58+
59+
def migrate_legacy_notification_subscriptions(
60+
dry_run=False,
61+
batch_size=1000,
62+
timeout_per_batch=300,
63+
default_frequency='none'
64+
):
65+
66+
logger.info('Starting legacy notification subscription migration...')
2767

2868
PROVIDER_BASED_LEGACY_NOTIFICATION_TYPES = [
29-
'new_pending_submissions', 'new_pending_withdraw_requests'
69+
'new_pending_submissions',
70+
'new_pending_withdraw_requests',
71+
'reviews_submission_confirmation',
72+
'reviews_moderator_submission_confirmation',
73+
'reviews_reject_confirmation',
74+
'reviews_accept_confirmation',
75+
'reviews_resubmission_confirmation',
76+
'reviews_comment_edited',
77+
'contributor_added_preprint',
78+
'confirm_email_moderation',
79+
'moderator_added',
80+
'confirm_email_preprints',
81+
'user_invite_preprint',
3082
]
83+
def timeout_handler(signum, frame):
84+
raise TimeoutError("Batch processing timed out")
85+
86+
# Precompute notification type IDs
87+
notiftype_map = dict(NotificationType.objects.values_list('name', 'id'))
88+
89+
# Cache existing keys only once
90+
existing_keys = set(
91+
(
92+
user_id,
93+
content_type_id,
94+
int(object_id) if object_id else None,
95+
notification_type_id,
96+
)
97+
for user_id, content_type_id, object_id, notification_type_id in
98+
NotificationSubscription.objects.all().values_list(
99+
"user_id", "content_type_id", "object_id", "notification_type_id"
100+
)
101+
)
102+
103+
total = NotificationSubscriptionLegacy.objects.count()
104+
created = 0
105+
skipped = 0
106+
last_id = 0
107+
start_time_total = time.time()
108+
109+
while True:
110+
batch_start_time = time.time()
111+
# Fetch the next chunk directly from DB
112+
batch = list(
113+
NotificationSubscriptionLegacy.objects
114+
.filter(id__gt=last_id)
115+
.order_by("id")[:batch_size]
116+
)
117+
if not batch:
118+
break # done
119+
subscriptions_to_create = []
120+
121+
signal.signal(signal.SIGALRM, timeout_handler)
122+
signal.alarm(timeout_per_batch)
31123

32-
for legacy in NotificationSubscriptionLegacy.objects.all():
33-
event_name = legacy.event_name
34-
if event_name in PROVIDER_BASED_LEGACY_NOTIFICATION_TYPES:
35-
subscribed_object = legacy.provider
36-
elif subscribed_object := legacy.node:
37-
pass
38-
elif subscribed_object := legacy.user:
39-
pass
40-
else:
41-
raise NotImplementedError(f'Invalid Notification id {event_name}')
42-
content_type = ContentType.objects.get_for_model(subscribed_object.__class__)
43-
44-
if notification_name := EVENT_NAME_TO_NOTIFICATION_TYPE[event_name]:
45-
subscription, _ = NotificationSubscription.objects.get_or_create(
46-
notification_type=NotificationType.objects.get(name=notification_name),
47-
user=legacy.user,
48-
content_type=content_type,
49-
object_id=subscribed_object.id,
50-
message_frequency=('weekly' if legacy.email_digest.exists() else 'none'),
124+
try:
125+
for legacy in batch:
126+
event_name = legacy.event_name
127+
if event_name in PROVIDER_BASED_LEGACY_NOTIFICATION_TYPES:
128+
subscribed_object = legacy.provider
129+
if not subscribed_object:
130+
skipped += 1
131+
continue
132+
elif legacy.node:
133+
subscribed_object = legacy.node
134+
elif legacy.user:
135+
subscribed_object = legacy.user
136+
else:
137+
skipped += 1
138+
continue
139+
140+
content_type = ContentType.objects.get_for_model(subscribed_object.__class__)
141+
notif_enum = EVENT_NAME_TO_NOTIFICATION_TYPE.get(event_name)
142+
if not notif_enum:
143+
skipped += 1
144+
continue
145+
146+
notification_type_id = notiftype_map.get(notif_enum)
147+
if not notification_type_id:
148+
skipped += 1
149+
continue
150+
151+
key = (
152+
legacy.user_id,
153+
content_type.id,
154+
int(subscribed_object.id),
155+
notification_type_id,
156+
)
157+
if key in existing_keys:
158+
skipped += 1
159+
continue
160+
161+
frequency = 'weekly' if getattr(legacy, 'email_digest', False) else default_frequency
162+
163+
if dry_run:
164+
created += 1
165+
else:
166+
subscriptions_to_create.append(NotificationSubscription(
167+
notification_type_id=notification_type_id,
168+
user_id=legacy.user_id,
169+
content_type=content_type,
170+
object_id=subscribed_object.id,
171+
message_frequency=frequency,
172+
))
173+
existing_keys.add(key)
174+
175+
if not dry_run and subscriptions_to_create:
176+
with transaction.atomic():
177+
NotificationSubscription.objects.bulk_create(
178+
subscriptions_to_create,
179+
ignore_conflicts=True,
180+
)
181+
created += len(subscriptions_to_create)
182+
183+
# Logging with ETA
184+
batch_time = time.time() - batch_start_time
185+
processed = last_id + len(batch)
186+
rate = processed / (time.time() - start_time_total)
187+
eta = (total - processed) / rate if rate else 0
188+
189+
logger.info(
190+
f"Processed batch {last_id}-{last_id + len(batch)} "
191+
f"in {batch_time:.2f}s | "
192+
f"Progress {processed}/{total} ({processed/total:.1%}) | "
193+
f"ETA ~ {eta/60:.1f} min"
51194
)
52-
logger.info(f'Created NotificationType "{event_name}" with content_type {content_type} with {subscription}')
195+
196+
except TimeoutError:
197+
logger.error(f"Batch {last_id}-{last_id + len(batch)} timed out, skipping.")
198+
skipped += len(batch)
199+
except Exception as e:
200+
logger.exception(f"Batch {last_id}-{last_id + len(batch)} failed: {e}")
201+
skipped += len(batch)
202+
finally:
203+
signal.alarm(0)
204+
last_id = batch[-1].id
205+
206+
elapsed_total = time.time() - start_time_total
207+
logger.info(
208+
f"Migration {'(dry-run) ' if dry_run else ''}completed in {elapsed_total:.2f} seconds: "
209+
f"{created} created, {skipped} skipped."
210+
)
211+
53212

54213
class Command(BaseCommand):
55214
help = 'Migrate legacy NotificationSubscriptionLegacy objects to new Notification app models.'
56215

216+
def add_arguments(self, parser):
217+
parser.add_argument(
218+
'--dry-run',
219+
action='store_true',
220+
help='Run migration in dry-run mode (no DB changes will be committed).'
221+
)
222+
57223
def handle(self, *args, **options):
58-
with transaction.atomic():
59-
populate_notification_types(args, options)
224+
dry_run = options['dry_run']
225+
226+
try:
227+
with time_limit(TIMEOUT_SECONDS):
228+
if not dry_run:
229+
with transaction.atomic():
230+
logger.info("Populating notification types...")
231+
populate_notification_types(args, options)
232+
233+
with transaction.atomic():
234+
migrate_legacy_notification_subscriptions(dry_run=dry_run)
235+
236+
except TimeoutError:
237+
logger.error("Migration timed out. Rolling back changes.")
238+
raise CommandError("Migration failed due to timeout")
239+
except Exception as e:
240+
logger.exception("Migration failed. Rolling back changes.")
241+
raise CommandError(str(e))
60242

61-
with transaction.atomic():
62-
migrate_legacy_notification_subscriptions(args, options)

0 commit comments

Comments
 (0)