diff --git a/zerver/lib/digest.py b/zerver/lib/digest.py index a6ade5cf69..d67d60e1ac 100644 --- a/zerver/lib/digest.py +++ b/zerver/lib/digest.py @@ -243,13 +243,17 @@ def bulk_get_digest_context(users: List[UserProfile], cutoff: float) -> Dict[int stream_map = get_stream_map(user_ids) + long_term_idle_user_ids = [user.id for user in users if user.long_term_idle] + + recently_modified_streams = get_modified_streams(long_term_idle_user_ids, cutoff_date) + all_stream_ids = set() for user in users: stream_ids = stream_map[user.id] if user.long_term_idle: - stream_ids -= streams_recently_modified_for_user(user, cutoff_date) + stream_ids -= recently_modified_streams.get(user.id, set()) all_stream_ids |= stream_ids @@ -341,18 +345,25 @@ def bulk_write_realm_audit_logs(users: List[UserProfile]) -> None: def handle_digest_email(user_id: int, cutoff: float) -> None: bulk_handle_digest_email([user_id], cutoff) -def streams_recently_modified_for_user(user: UserProfile, cutoff_date: datetime.datetime) -> Set[int]: +def get_modified_streams(user_ids: List[int], cutoff_date: datetime.datetime) -> Dict[int, Set[int]]: events = [ RealmAuditLog.SUBSCRIPTION_CREATED, RealmAuditLog.SUBSCRIPTION_ACTIVATED, RealmAuditLog.SUBSCRIPTION_DEACTIVATED, ] - # Streams where the user's subscription was changed - modified_streams = RealmAuditLog.objects.filter( - realm=user.realm, - modified_user=user, + # Get rows where the users' subscriptions have changed. + rows = RealmAuditLog.objects.filter( + modified_user_id__in=user_ids, event_time__gt=cutoff_date, - event_type__in=events).values_list('modified_stream_id', flat=True) + event_type__in=events, + ).values("modified_user_id", "modified_stream_id").distinct() - return set(modified_streams) + result: Dict[int, Set[int]] = defaultdict(set) + + for row in rows: + user_id = row["modified_user_id"] + stream_id = row["modified_stream_id"] + result[user_id].add(stream_id) + + return result diff --git a/zerver/tests/test_digest.py b/zerver/tests/test_digest.py index 0a2059e61c..9f4142800a 100644 --- a/zerver/tests/test_digest.py +++ b/zerver/tests/test_digest.py @@ -14,8 +14,8 @@ from zerver.lib.digest import ( bulk_write_realm_audit_logs, enqueue_emails, gather_new_streams, + get_modified_streams, handle_digest_email, - streams_recently_modified_for_user, ) from zerver.lib.message import get_last_message_id from zerver.lib.streams import create_stream_if_needed @@ -180,7 +180,7 @@ class TestDigestEmailMessages(ZulipTestCase): with cache_tries_captured() as cache_tries: bulk_handle_digest_email(digest_user_ids, cutoff) - self.assert_length(queries, 37) + self.assert_length(queries, 34) self.assert_length(cache_tries, 0) self.assertEqual(mock_send_future_email.call_count, len(digest_users)) @@ -235,39 +235,31 @@ class TestDigestEmailMessages(ZulipTestCase): self.unsubscribe(othello, 'Denmark') self.subscribe(othello, 'Denmark') - self.assertEqual( - streams_recently_modified_for_user(othello, one_hour_ago), - {denmark.id} - ) + recent_streams = get_modified_streams([othello.id], one_hour_ago) + self.assertEqual(recent_streams[othello.id], {denmark.id}) # Backdate all our logs (so that Denmark will no longer # appear like a recently modified stream for Othello). RealmAuditLog.objects.all().update(event_time=two_hours_ago) # Now Denmark no longer appears recent to Othello. - self.assertEqual( - streams_recently_modified_for_user(othello, one_hour_ago), - set() - ) + recent_streams = get_modified_streams([othello.id], one_hour_ago) + self.assertEqual(recent_streams[othello.id], set()) # Unsubscribe and subscribe from a stream self.unsubscribe(othello, 'Verona') self.subscribe(othello, 'Verona') # Now, Verona, but not Denmark, appears recent. - self.assertEqual( - streams_recently_modified_for_user(othello, one_hour_ago), - {verona.id}, - ) + recent_streams = get_modified_streams([othello.id], one_hour_ago) + self.assertEqual(recent_streams[othello.id], {verona.id}) # make sure we don't mix up Othello and Cordelia self.unsubscribe(cordelia, 'Denmark') self.subscribe(cordelia, 'Denmark') - self.assertEqual( - streams_recently_modified_for_user(cordelia, one_hour_ago), - {denmark.id} - ) + recent_streams = get_modified_streams([othello.id, cordelia.id], one_hour_ago) + self.assertEqual(recent_streams[cordelia.id], {denmark.id}) def active_human_users(self, realm: Realm) -> List[UserProfile]: users = list(UserProfile.objects.filter(