stream_traffic: Update get_streams_traffic to return None for zephyr realm.

Instead of having a "realm.is_zephyr_mirror_realm" check for every
get_streams_traffic call, this commit udpates get_streams_traffic to
accept realm as parameter and return "None" for zephyr mirror realm.
This commit is contained in:
Sahil Batra 2023-08-07 20:17:48 +05:30 committed by Tim Abbott
parent 7b256f019a
commit 6776e380b2
5 changed files with 16 additions and 23 deletions

View File

@ -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)
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

View File

@ -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)

View File

@ -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)
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"])

View File

@ -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.

View File

@ -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)})