From faba34dae4661563368f8be6002e4903129c23cd Mon Sep 17 00:00:00 2001 From: Steve Howell Date: Sun, 29 Oct 2017 12:19:57 -0700 Subject: [PATCH] Simplify bulk_remove_subscriptions(). We extract get_bulk_stream_subscriber_info() from this function to remove some of the complexity. Also, in that new function we avoid a hop to the database by querying on stream ids instead of recipient ids. The query that gets changed here does require a join to the recipient table (to get the stream id), so it's a little bit of a tradeoff. --- zerver/lib/actions.py | 64 ++++++++++++++++++++----------- zerver/lib/stream_subscription.py | 34 +++++++++++++++- zerver/tests/test_subs.py | 6 +-- 3 files changed, 78 insertions(+), 26 deletions(-) diff --git a/zerver/lib/actions.py b/zerver/lib/actions.py index 8998eb314b..e47fb6973d 100644 --- a/zerver/lib/actions.py +++ b/zerver/lib/actions.py @@ -37,6 +37,7 @@ from zerver.lib.send_email import send_email, FromAddress from zerver.lib.stream_subscription import ( get_active_subscriptions_for_stream_id, get_active_subscriptions_for_stream_ids, + get_bulk_stream_subscriber_info, get_stream_subscriptions_for_user, get_stream_subscriptions_for_users, num_subscribers_for_stream_id, @@ -2274,28 +2275,44 @@ def bulk_remove_subscriptions(users, streams, acting_user=None): # type: (Iterable[UserProfile], Iterable[Stream], Optional[UserProfile]) -> Tuple[List[Tuple[UserProfile, Stream]], List[Tuple[UserProfile, Stream]]] users = list(users) + streams = list(streams) - recipients_map = bulk_get_recipients(Recipient.STREAM, - [stream.id for stream in streams]) # type: Mapping[int, Recipient] - stream_map = {} # type: Dict[int, Stream] - for stream in streams: - stream_map[recipients_map[stream.id].id] = stream + stream_dict = {stream.id: stream for stream in streams} - subs_by_user = dict((user_profile.id, []) for user_profile in users) # type: Dict[int, List[Subscription]] - for sub in Subscription.objects.select_related("user_profile").filter(user_profile__in=users, - recipient__in=list(recipients_map.values()), - active=True): - subs_by_user[sub.user_profile_id].append(sub) + existing_subs_by_user = get_bulk_stream_subscriber_info(users, stream_dict) + + def get_non_subscribed_tups(): + # type: () -> List[Tuple[UserProfile, Stream]] + stream_ids = {stream.id for stream in streams} + + not_subscribed = [] # type: List[Tuple[UserProfile, Stream]] + + for user_profile in users: + user_sub_stream_info = existing_subs_by_user[user_profile.id] + + subscribed_stream_ids = { + stream.id + for (sub, stream) in user_sub_stream_info + } + not_subscribed_stream_ids = stream_ids - subscribed_stream_ids + + for stream_id in not_subscribed_stream_ids: + stream = stream_dict[stream_id] + not_subscribed.append((user_profile, stream)) + + return not_subscribed + + not_subscribed = get_non_subscribed_tups() subs_to_deactivate = [] # type: List[Tuple[Subscription, Stream]] - not_subscribed = [] # type: List[Tuple[UserProfile, Stream]] - for user_profile in users: - recipients_to_unsub = set([recipient.id for recipient in recipients_map.values()]) - for sub in subs_by_user[user_profile.id]: - recipients_to_unsub.remove(sub.recipient_id) - subs_to_deactivate.append((sub, stream_map[sub.recipient_id])) - for recipient_id in recipients_to_unsub: - not_subscribed.append((user_profile, stream_map[recipient_id])) + sub_ids_to_deactivate = [] # type: List[int] + + # This loop just flattens out our data into big lists for + # bulk operations. + for tup_list in existing_subs_by_user.values(): + for (sub, stream) in tup_list: + subs_to_deactivate.append((sub, stream)) + sub_ids_to_deactivate.append(sub.id) our_realm = users[0].realm @@ -2303,8 +2320,9 @@ def bulk_remove_subscriptions(users, streams, acting_user=None): # transaction isolation level. with transaction.atomic(): occupied_streams_before = list(get_occupied_streams(our_realm)) - Subscription.objects.filter(id__in=[sub.id for (sub, stream_name) in - subs_to_deactivate]).update(active=False) + Subscription.objects.filter( + id__in=sub_ids_to_deactivate, + ) .update(active=False) occupied_streams_after = list(get_occupied_streams(our_realm)) # Log Subscription Activities in RealmAuditLog @@ -2373,8 +2391,10 @@ def bulk_remove_subscriptions(users, streams, acting_user=None): user_id=removed_user.id) send_event(event, peer_user_ids) - return ([(sub.user_profile, stream) for (sub, stream) in subs_to_deactivate], - not_subscribed) + return ( + [(sub.user_profile, stream) for (sub, stream) in subs_to_deactivate], + not_subscribed, + ) def log_subscription_property_change(user_email, stream_name, property, value): # type: (Text, Text, Text, Any) -> None diff --git a/zerver/lib/stream_subscription.py b/zerver/lib/stream_subscription.py index 47ce7c73d2..0cdeafcaab 100644 --- a/zerver/lib/stream_subscription.py +++ b/zerver/lib/stream_subscription.py @@ -1,8 +1,10 @@ -from typing import List +from typing import Dict, List, Tuple +from mypy_extensions import TypedDict from django.db.models.query import QuerySet from zerver.models import ( Recipient, + Stream, Subscription, UserProfile, ) @@ -37,6 +39,36 @@ def get_stream_subscriptions_for_users(user_profiles): recipient__type=Recipient.STREAM, ) +SubInfo = TypedDict('SubInfo', { + 'sub': Subscription, + 'stream': Stream, +}) + +def get_bulk_stream_subscriber_info(user_profiles, stream_dict): + # type: (List[UserProfile], Dict[int, Stream]) -> Dict[int, List[Tuple[Subscription, Stream]]] + + stream_ids = stream_dict.keys() + + result = { + user_profile.id: [] + for user_profile in user_profiles + } # type: Dict[int, List[Tuple[Subscription, Stream]]] + + subs = Subscription.objects.filter( + user_profile__in=user_profiles, + recipient__type=Recipient.STREAM, + recipient__type_id__in=stream_ids, + active=True, + ).select_related('user_profile', 'recipient') + + for sub in subs: + user_profile_id = sub.user_profile_id + stream_id = sub.recipient.type_id + stream = stream_dict[stream_id] + result[user_profile_id].append((sub, stream)) + + return result + def num_subscribers_for_stream_id(stream_id): # type: (int) -> int return get_active_subscriptions_for_stream_id(stream_id).filter( diff --git a/zerver/tests/test_subs.py b/zerver/tests/test_subs.py index 9082c35fc8..293798dffb 100644 --- a/zerver/tests/test_subs.py +++ b/zerver/tests/test_subs.py @@ -630,7 +630,7 @@ class StreamAdminTest(ZulipTestCase): those you aren't on. """ result = self.attempt_unsubscribe_of_principal( - query_count=15, is_admin=True, is_subbed=True, invite_only=False, + query_count=14, is_admin=True, is_subbed=True, invite_only=False, other_user_subbed=True) json = self.assert_json_success(result) self.assertEqual(len(json["removed"]), 1) @@ -643,7 +643,7 @@ class StreamAdminTest(ZulipTestCase): are on. """ result = self.attempt_unsubscribe_of_principal( - query_count=15, is_admin=True, is_subbed=True, invite_only=True, + query_count=14, is_admin=True, is_subbed=True, invite_only=True, other_user_subbed=True) json = self.assert_json_success(result) self.assertEqual(len(json["removed"]), 1) @@ -716,7 +716,7 @@ class StreamAdminTest(ZulipTestCase): fails gracefully. """ result = self.attempt_unsubscribe_of_principal( - query_count=12, is_admin=True, is_subbed=False, invite_only=False, + query_count=11, is_admin=True, is_subbed=False, invite_only=False, other_user_subbed=False) json = self.assert_json_success(result) self.assertEqual(len(json["removed"]), 0)