zulip/zerver/lib/stream_traffic.py

import datetime
from typing import Dict, Optional, Set

from django.db.models import Sum
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], 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)

    query = StreamCount.objects.filter(property=stat.property, end_time__gt=traffic_from)
    query = query.filter(stream_id__in=stream_ids)

    traffic_list = query.values("stream_id").annotate(value=Sum("value"))
    traffic_dict = {}
    for traffic in traffic_list:
        traffic_dict[traffic["stream_id"]] = traffic["value"]

    return traffic_dict


def round_to_2_significant_digits(number: int) -> int:
    return int(round(number, 2 - len(str(number))))


STREAM_TRAFFIC_CALCULATION_MIN_AGE_DAYS = 7


def get_average_weekly_stream_traffic(
    stream_id: int, stream_date_created: datetime.datetime, recent_traffic: Dict[int, int]
) -> Optional[int]:
    try:
        stream_traffic = recent_traffic[stream_id]
    except KeyError:
        stream_traffic = 0

    stream_age = (timezone_now() - stream_date_created).days

    if stream_age >= 28:
        average_weekly_traffic = int(stream_traffic // 4)
    elif stream_age >= STREAM_TRAFFIC_CALCULATION_MIN_AGE_DAYS:
        average_weekly_traffic = int(stream_traffic * 7 // stream_age)
    else:
        return None

    if average_weekly_traffic == 0 and stream_traffic > 0:
        average_weekly_traffic = 1

    return round_to_2_significant_digits(average_weekly_traffic)