From 56949b47885a77348bb49557cdd49f1bd5193826 Mon Sep 17 00:00:00 2001 From: Tim Abbott Date: Fri, 10 May 2013 11:43:27 -0400 Subject: [PATCH] Create subscription objects in bulk by user_profile. This, combined with acrefoot's work on sending the notifications in bulk, resolves trac #1142 -- we do only 10 database queries and the whole operation completes in about 300ms on my laptop. (imported from commit 36b5bb836bc6c713903d1ca72e39af87775dc469) --- zephyr/lib/actions.py | 82 +++++++++++++++++++++++++++++++++---------- zephyr/views.py | 13 ++++--- 2 files changed, 70 insertions(+), 25 deletions(-) diff --git a/zephyr/lib/actions.py b/zephyr/lib/actions.py index 9c75507697..62ab94cc9c 100644 --- a/zephyr/lib/actions.py +++ b/zephyr/lib/actions.py @@ -415,10 +415,13 @@ def internal_send_message(sender_email, recipient_type_name, recipients, if ret is not None: logging.error("Error sending internal message by %s: %s" % (sender_email, ret)) -def get_stream_colors(user_profile): - return [(sub["name"], sub["color"]) for sub in gather_subscriptions(user_profile)] - def pick_color(user_profile): + subs = Subscription.objects.filter(user_profile=user_profile, + active=True, + recipient__type=Recipient.STREAM) + return pick_color_helper(user_profile, subs) + +def pick_color_helper(user_profile, subs): # These colors are shared with the palette in subs.js. stream_assignment_colors = [ "#76ce90", "#fae589", "#a6c7e5", "#e79ab5", @@ -427,7 +430,7 @@ def pick_color(user_profile): "#ee7e4a", "#a6dcbf", "#95a5fd", "#53a063", "#9987e1", "#e4523d", "#c2c2c2", "#4f8de4", "#c6a8ad", "#e7cc4d", "#c8bebf", "#a47462"] - used_colors = [elt[1] for elt in get_stream_colors(user_profile) if elt[1]] + used_colors = [sub.color for sub in subs if sub.active] available_colors = filter(lambda x: x not in used_colors, stream_assignment_colors) @@ -450,6 +453,62 @@ def set_stream_color(user_profile, stream_name, color=None): subscription.save(update_fields=["color"]) return color +def notify_new_subscription(user_profile, stream, subscription, no_log=False): + if not no_log: + log_event({'type': 'subscription_added', + 'user': user_profile.email, + 'name': stream.name, + 'domain': stream.realm.domain}) + + notice = dict(event=dict(type="subscription", op="add", + subscription=dict(name=stream.name, + in_home_view=subscription.in_home_view, + invite_only=stream.invite_only, + color=subscription.color)), + users=[user_profile.id]) + tornado_callbacks.send_notification(notice) + +# This function assumes that the users are known to not be subscribers +# of the stream (e.g. because the stream was created by this same query) +def bulk_add_subscriptions(stream, users): + recipient = get_recipient(Recipient.STREAM, stream.id) + all_subs = Subscription.objects.filter(user_profile__in=users, + recipient__type=Recipient.STREAM) + + subs_by_user = defaultdict(list) + for sub in all_subs: + subs_by_user[sub.user_profile_id].append(sub) + + already_subscribed = [] + subs_to_activate = [] + users_needing_new_subs = [] + for user_profile in users: + needs_new_sub = True + for sub in subs_by_user[user_profile.id]: + if sub.recipient_id == recipient.id: + needs_new_sub = False + if sub.active: + already_subscribed.append(user_profile) + else: + subs_to_activate.append(sub) + if needs_new_sub: + users_needing_new_subs.append(user_profile) + + subs_to_add = [] + for user_profile in users_needing_new_subs: + color = pick_color_helper(user_profile, subs_by_user[user_profile.id]) + subs_to_add.append(Subscription(user_profile=user_profile, + active=True, color=color, + recipient=recipient)) + Subscription.objects.bulk_create(subs_to_add) + Subscription.objects.filter(id__in=[s.id for s in subs_to_activate]).update(active=True) + + for sub in subs_to_add + subs_to_activate: + notify_new_subscription(sub.user_profile, stream, sub) + return (users_needing_new_subs + [sub.user_profile for sub in subs_to_activate], + already_subscribed) + +# When changing this, also change bulk_add_subscriptions def do_add_subscription(user_profile, stream, no_log=False): recipient = get_recipient(Recipient.STREAM, stream.id) color = pick_color(user_profile) @@ -462,20 +521,7 @@ def do_add_subscription(user_profile, stream, no_log=False): subscription.active = True subscription.save(update_fields=["active"]) if did_subscribe: - if not no_log: - log_event({'type': 'subscription_added', - 'user': user_profile.email, - 'name': stream.name, - 'domain': stream.realm.domain}) - - notice = dict(event=dict(type="subscription", op="add", - subscription=dict(name=stream.name, - in_home_view=subscription.in_home_view, - invite_only=stream.invite_only, - color=subscription.color)), - users=[user_profile.id]) - tornado_callbacks.send_notification(notice) - + notify_new_subscription(user_profile, stream, subscription, no_log) return did_subscribe def do_remove_subscription(user_profile, stream, no_log=False): diff --git a/zephyr/views.py b/zephyr/views.py index 1b23f1444b..75e1d5d2e2 100644 --- a/zephyr/views.py +++ b/zephyr/views.py @@ -27,7 +27,7 @@ from zephyr.lib.actions import do_add_subscription, do_remove_subscription, \ do_send_confirmation_email, do_activate_user, do_create_user, check_send_message, \ log_subscription_property_change, internal_send_message, \ create_stream_if_needed, gather_subscriptions, subscribed_to_stream, \ - update_user_presence, set_stream_color, get_stream_colors, update_message_flags, \ + update_user_presence, bulk_add_subscriptions, update_message_flags, \ recipient_for_emails, extract_recipients, do_events_register, do_finish_tutorial, \ get_status_dict, do_change_enable_offline_email_notifications, \ do_update_onboarding_steps @@ -1175,12 +1175,11 @@ def add_subscriptions_backend(request, user_profile, result = dict(subscribed=defaultdict(list), already_subscribed=defaultdict(list)) for stream in streams: - for subscriber in subscribers: - did_subscribe = do_add_subscription(subscriber, stream) - if did_subscribe: - result["subscribed"][subscriber.email].append(stream.name) - else: - result["already_subscribed"][subscriber.email].append(stream.name) + (subscribed, already_subscribed) = bulk_add_subscriptions(stream, subscribers) + for subscriber in subscribed: + result["subscribed"][subscriber.email].append(stream.name) + for subscriber in already_subscribed: + result["already_subscribed"][subscriber.email].append(stream.name) private_streams[stream.name] = stream.invite_only # Inform the user if someone else subscribed them to stuff