diff --git a/zerver/actions/streams.py b/zerver/actions/streams.py index 58361618c2..292c94df70 100644 --- a/zerver/actions/streams.py +++ b/zerver/actions/streams.py @@ -210,10 +210,7 @@ def do_reactivate_stream( event_time=timezone_now(), ) - recent_traffic = None - if not realm.is_zephyr_mirror_realm: - # We do not need stream traffic data for streams in zephyr mirroring realm. - recent_traffic = get_streams_traffic({stream.id}) + recent_traffic = get_streams_traffic({stream.id}, realm) # All admins always get to know about private streams' existence, # but we only subscribe the realm owners. @@ -443,11 +440,8 @@ def send_stream_creation_events_for_previously_inaccessible_streams( altered_user_dict: Dict[int, Set[int]], altered_guests: Set[int], ) -> None: - recent_traffic = None - if not realm.is_zephyr_mirror_realm: - # We do not need stream traffic data for streams in zephyr mirroring realm. - stream_ids = set(altered_user_dict.keys()) - recent_traffic = get_streams_traffic(stream_ids) + stream_ids = set(altered_user_dict.keys()) + recent_traffic = get_streams_traffic(stream_ids, realm) for stream_id, stream_users_ids in altered_user_dict.items(): stream = stream_dict[stream_id] @@ -1046,10 +1040,7 @@ def do_change_stream_permission( non_guest_user_ids = set(active_non_guest_user_ids(stream.realm_id)) notify_stream_creation_ids = non_guest_user_ids - old_can_access_stream_user_ids - recent_traffic = None - if not realm.is_zephyr_mirror_realm: - # We do not need stream traffic data for streams in zephyr mirroing realm. - recent_traffic = get_streams_traffic({stream.id}) + recent_traffic = get_streams_traffic({stream.id}, realm) send_stream_creation_event(realm, stream, list(notify_stream_creation_ids), recent_traffic) # Add subscribers info to the stream object. We need to send peer_add diff --git a/zerver/lib/stream_traffic.py b/zerver/lib/stream_traffic.py index 4a08812c20..b0d71f17a1 100644 --- a/zerver/lib/stream_traffic.py +++ b/zerver/lib/stream_traffic.py @@ -6,9 +6,16 @@ from django.utils.timezone import now as timezone_now from analytics.lib.counts import COUNT_STATS from analytics.models import StreamCount +from zerver.models import Realm -def get_streams_traffic(stream_ids: Set[int]) -> Dict[int, int]: +def get_streams_traffic( + stream_ids: Set[int], realm: Optional[Realm] = None +) -> Optional[Dict[int, int]]: + if realm is not None and realm.is_zephyr_mirror_realm: + # We do not need traffic data for streams in zephyr mirroring realm. + return None + stat = COUNT_STATS["messages_in_stream:is_bot:day"] traffic_from = timezone_now() - datetime.timedelta(days=28) diff --git a/zerver/lib/streams.py b/zerver/lib/streams.py index e9a2e9375b..2d5f92e2c9 100644 --- a/zerver/lib/streams.py +++ b/zerver/lib/streams.py @@ -941,11 +941,8 @@ def do_get_streams( # Don't bother going to the database with no valid sources return [] - recent_traffic = None - if not user_profile.realm.is_zephyr_mirror_realm: - # We do not need stream traffic data for streams in zephyr mirroring realm. - stream_ids = {stream.id for stream in streams} - recent_traffic = get_streams_traffic(stream_ids) + stream_ids = {stream.id for stream in streams} + recent_traffic = get_streams_traffic(stream_ids, user_profile.realm) stream_dicts = [stream_to_dict(stream, recent_traffic) for stream in streams] stream_dicts.sort(key=lambda elt: elt["name"]) diff --git a/zerver/lib/subscription_info.py b/zerver/lib/subscription_info.py index f718c1ae47..df1346eecf 100644 --- a/zerver/lib/subscription_info.py +++ b/zerver/lib/subscription_info.py @@ -431,6 +431,7 @@ def gather_subscriptions_helper( traffic_stream_ids = {get_stream_id(sub_dict) for sub_dict in sub_dicts} recent_traffic = get_streams_traffic(stream_ids=traffic_stream_ids) + assert recent_traffic is not None # Okay, now we finally get to populating our main results, which # will be these three lists. diff --git a/zerver/views/streams.py b/zerver/views/streams.py index cb40001dab..8eef598843 100644 --- a/zerver/views/streams.py +++ b/zerver/views/streams.py @@ -895,10 +895,7 @@ def get_stream_backend( ) -> HttpResponse: (stream, sub) = access_stream_by_id(user_profile, stream_id, allow_realm_admin=True) - recent_traffic = None - if not user_profile.realm.is_zephyr_mirror_realm: - # We do not need stream traffic data in zephyr mirroring realm. - recent_traffic = get_streams_traffic({stream.id}) + recent_traffic = get_streams_traffic({stream.id}, user_profile.realm) return json_success(request, data={"stream": stream_to_dict(stream, recent_traffic)})