diff --git a/analytics/lib/counts.py b/analytics/lib/counts.py index 373b78b6e5..18d18663c9 100644 --- a/analytics/lib/counts.py +++ b/analytics/lib/counts.py @@ -23,6 +23,9 @@ from zerver.lib.logging_util import log_to_file from zerver.lib.timestamp import ceiling_to_day, ceiling_to_hour, floor_to_hour, verify_UTC from zerver.models import Message, Realm, RealmAuditLog, Stream, UserActivityInterval, UserProfile +if settings.ZILENCER_ENABLED: + from zilencer.models import RemoteInstallationCount, RemoteZulipServer + ## Logging setup ## logger = logging.getLogger("zulip.management") @@ -293,7 +296,7 @@ def do_aggregate_to_summary_table( # called from zerver.actions; should not throw any errors def do_increment_logging_stat( - model_object_for_bucket: Union[Realm, UserProfile, Stream], + model_object_for_bucket: Union[Realm, UserProfile, Stream, "RemoteZulipServer"], stat: CountStat, subgroup: Optional[Union[str, int, bool]], event_time: datetime, @@ -305,13 +308,20 @@ def do_increment_logging_stat( table = stat.data_collector.output_table if table == RealmCount: assert isinstance(model_object_for_bucket, Realm) - id_args: Dict[str, Union[Realm, UserProfile, Stream]] = {"realm": model_object_for_bucket} + id_args: Dict[str, Optional[Union[Realm, UserProfile, Stream, "RemoteZulipServer"]]] = { + "realm": model_object_for_bucket + } elif table == UserCount: assert isinstance(model_object_for_bucket, UserProfile) id_args = {"realm": model_object_for_bucket.realm, "user": model_object_for_bucket} - else: # StreamCount + elif table == StreamCount: assert isinstance(model_object_for_bucket, Stream) id_args = {"realm": model_object_for_bucket.realm, "stream": model_object_for_bucket} + elif table == RemoteInstallationCount: + assert isinstance(model_object_for_bucket, RemoteZulipServer) + id_args = {"server": model_object_for_bucket, "remote_id": None} + else: + raise AssertionError("Unsupported CountStat output_table") if stat.frequency == CountStat.DAY: end_time = ceiling_to_day(event_time) @@ -833,6 +843,15 @@ def get_count_stats(realm: Optional[Realm] = None) -> Dict[str, CountStat]: ), ] + if settings.ZILENCER_ENABLED: + count_stats_.append( + LoggingCountStat( + "mobile_pushes_received::day", + RemoteInstallationCount, + CountStat.DAY, + ) + ) + return OrderedDict((stat.property, stat) for stat in count_stats_) diff --git a/analytics/tests/test_counts.py b/analytics/tests/test_counts.py index ae5e6219cc..1995d5a7c0 100644 --- a/analytics/tests/test_counts.py +++ b/analytics/tests/test_counts.py @@ -3,9 +3,11 @@ from typing import Any, Dict, List, Optional, Tuple, Type from unittest import mock import orjson +import time_machine from django.apps import apps from django.db import models from django.db.models import Sum +from django.test import override_settings from django.utils.timezone import now as timezone_now from psycopg2.sql import SQL, Literal from typing_extensions import override @@ -53,14 +55,20 @@ from zerver.actions.user_activity import update_user_activity_interval from zerver.actions.users import do_deactivate_user from zerver.lib.create_user import create_user from zerver.lib.exceptions import InvitationError +from zerver.lib.push_notifications import ( + get_message_payload_apns, + get_message_payload_gcm, + hex_to_b64, +) from zerver.lib.test_classes import ZulipTestCase -from zerver.lib.timestamp import TimeZoneNotUTCError, floor_to_day +from zerver.lib.timestamp import TimeZoneNotUTCError, ceiling_to_day, floor_to_day from zerver.lib.topic import DB_TOPIC_NAME from zerver.lib.utils import assert_is_not_none from zerver.models import ( Client, Huddle, Message, + NotificationTriggers, PreregistrationUser, Realm, RealmAuditLog, @@ -74,7 +82,7 @@ from zerver.models import ( get_user, is_cross_realm_bot_email, ) -from zilencer.models import RemoteInstallationCount, RemoteZulipServer +from zilencer.models import RemoteInstallationCount, RemotePushDeviceToken, RemoteZulipServer from zilencer.views import get_last_id_from_server @@ -236,7 +244,7 @@ class AnalyticsTestCase(ZulipTestCase): kwargs[arg_keys[i]] = values[i] for key, value in defaults.items(): kwargs[key] = kwargs.get(key, value) - if table is not InstallationCount and "realm" not in kwargs: + if table not in [InstallationCount, RemoteInstallationCount] and "realm" not in kwargs: if "user" in kwargs: kwargs["realm"] = kwargs["user"].realm elif "stream" in kwargs: @@ -1377,6 +1385,86 @@ class TestLoggingCountStats(AnalyticsTestCase): ], ) + @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com") + def test_mobile_pushes_received_count(self) -> None: + self.server_uuid = "6cde5f7a-1f7e-4978-9716-49f69ebfc9fe" + self.server = RemoteZulipServer.objects.create( + uuid=self.server_uuid, + api_key="magic_secret_api_key", + hostname="demo.example.com", + last_updated=timezone_now(), + ) + + hamlet = self.example_user("hamlet") + token = "aaaa" + + RemotePushDeviceToken.objects.create( + kind=RemotePushDeviceToken.GCM, + token=hex_to_b64(token), + user_uuid=(hamlet.uuid), + server=self.server, + ) + RemotePushDeviceToken.objects.create( + kind=RemotePushDeviceToken.GCM, + token=hex_to_b64(token + "aa"), + user_uuid=(hamlet.uuid), + server=self.server, + ) + RemotePushDeviceToken.objects.create( + kind=RemotePushDeviceToken.APNS, + token=hex_to_b64(token), + user_uuid=str(hamlet.uuid), + server=self.server, + ) + + message = Message( + sender=hamlet, + recipient=self.example_user("othello").recipient, + realm_id=hamlet.realm_id, + content="This is test content", + rendered_content="This is test content", + date_sent=timezone_now(), + sending_client=get_client("test"), + ) + message.set_topic_name("Test topic") + message.save() + gcm_payload, gcm_options = get_message_payload_gcm(hamlet, message) + apns_payload = get_message_payload_apns( + hamlet, message, NotificationTriggers.DIRECT_MESSAGE + ) + + payload = { + "user_id": hamlet.id, + "user_uuid": str(hamlet.uuid), + "gcm_payload": gcm_payload, + "apns_payload": apns_payload, + "gcm_options": gcm_options, + } + now = timezone_now() + with time_machine.travel(now, tick=False), mock.patch( + "zilencer.views.send_android_push_notification" + ), mock.patch("zilencer.views.send_apple_push_notification"), self.assertLogs( + "zilencer.views", level="INFO" + ): + result = self.uuid_post( + self.server_uuid, + "/api/v1/remotes/push/notify", + payload, + content_type="application/json", + subdomain="", + ) + self.assert_json_success(result) + + # There are 3 devices we created for the user, and the Count increment should + # match that number. + self.assertTableState( + RemoteInstallationCount, + ["property", "value", "subgroup", "server", "remote_id", "end_time"], + [ + ["mobile_pushes_received::day", 3, None, self.server, None, ceiling_to_day(now)], + ], + ) + def test_invites_sent(self) -> None: property = "invites_sent::day" diff --git a/zilencer/models.py b/zilencer/models.py index 0e497df05c..cd97caf5b8 100644 --- a/zilencer/models.py +++ b/zilencer/models.py @@ -135,6 +135,9 @@ class BaseRemoteCount(BaseCount): # The remote_id field is the id value of the corresponding *Count object # on the remote server. # It lets us deduplicate data from the remote server. + # Note: Some counts don't come from the remote server, but rather + # are stats we track on the bouncer server itself, pertaining to the remote server. + # E.g. mobile_pushes_received::day. Such counts will set this field to None. remote_id = models.IntegerField(null=True) class Meta: diff --git a/zilencer/views.py b/zilencer/views.py index 03842fc0e8..998b15a121 100644 --- a/zilencer/views.py +++ b/zilencer/views.py @@ -17,7 +17,7 @@ from django.utils.translation import gettext as err_ from django.views.decorators.csrf import csrf_exempt from pydantic import BaseModel, ConfigDict -from analytics.lib.counts import COUNT_STATS +from analytics.lib.counts import COUNT_STATS, do_increment_logging_stat from corporate.lib.stripe import do_deactivate_remote_server from zerver.decorator import require_post from zerver.lib.exceptions import JsonableError @@ -397,6 +397,13 @@ def remote_server_notify_push( len(android_devices), len(apple_devices), ) + do_increment_logging_stat( + server, + COUNT_STATS["mobile_pushes_received::day"], + None, + timezone_now(), + increment=len(android_devices) + len(apple_devices), + ) # Truncate incoming pushes to 200, due to APNs maximum message # sizes; see handle_remove_push_notification for the version of