From a67dd6dc1fb4b9946ff50fb57e08eaf99b3cf6cf Mon Sep 17 00:00:00 2001 From: Mateusz Mandera Date: Fri, 1 Dec 2023 14:52:44 +0100 Subject: [PATCH] realms: Call send_realms_only_to_push_bouncer at realm creation/import. --- zerver/actions/create_realm.py | 3 +++ zerver/lib/import_realm.py | 3 +++ zerver/lib/remote_server.py | 12 ++++++++++++ zerver/tests/test_import_export.py | 21 +++++++++++++++++++++ zerver/tests/test_realm.py | 17 +++++++++++++++++ zerver/worker/queue_processors.py | 12 +++++++++++- 6 files changed, 67 insertions(+), 1 deletion(-) diff --git a/zerver/actions/create_realm.py b/zerver/actions/create_realm.py index b3a447cba9..e1c9d32afa 100644 --- a/zerver/actions/create_realm.py +++ b/zerver/actions/create_realm.py @@ -14,6 +14,7 @@ from zerver.actions.realm_settings import ( do_deactivate_realm, ) from zerver.lib.bulk_create import create_users +from zerver.lib.remote_server import enqueue_register_realm_with_push_bouncer_if_needed from zerver.lib.server_initialization import create_internal_realm, server_initialized from zerver.lib.streams import ensure_stream, get_signups_stream from zerver.lib.user_groups import ( @@ -265,6 +266,8 @@ def do_create_realm( ] ) + enqueue_register_realm_with_push_bouncer_if_needed(realm) + # Create stream once Realm object has been saved notifications_stream = ensure_stream( realm, diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index d4ab904ef7..5fff552aa1 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -27,6 +27,7 @@ from zerver.lib.export import DATE_FIELDS, Field, Path, Record, TableData, Table from zerver.lib.markdown import markdown_convert from zerver.lib.markdown import version as markdown_version from zerver.lib.message import get_last_message_id +from zerver.lib.remote_server import enqueue_register_realm_with_push_bouncer_if_needed from zerver.lib.server_initialization import create_internal_realm, server_initialized from zerver.lib.streams import render_stream_description from zerver.lib.timestamp import datetime_to_timestamp @@ -1018,6 +1019,8 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea if "zerver_usergroup" not in data: set_default_for_realm_permission_group_settings(realm) + enqueue_register_realm_with_push_bouncer_if_needed(realm) + # Remap the user IDs for notification_bot and friends to their # appropriate IDs on this server internal_realm = get_realm(settings.SYSTEM_BOT_REALM) diff --git a/zerver/lib/remote_server.py b/zerver/lib/remote_server.py index 5ee884efca..245d8cc663 100644 --- a/zerver/lib/remote_server.py +++ b/zerver/lib/remote_server.py @@ -14,6 +14,7 @@ from version import ZULIP_VERSION from zerver.lib.exceptions import JsonableError, MissingRemoteRealmError from zerver.lib.export import floatify_datetime_fields from zerver.lib.outgoing_http import OutgoingSession +from zerver.lib.queue import queue_json_publish from zerver.models import OrgTypeEnum, Realm, RealmAuditLog @@ -299,3 +300,14 @@ def send_realms_only_to_push_bouncer() -> None: # We don't catch JsonableError here, because we want it to propagate further # to either explicitly, loudly fail or be error-handled by the caller. send_to_push_bouncer("POST", "server/analytics", request) + + +def enqueue_register_realm_with_push_bouncer_if_needed(realm: Realm) -> None: + from zerver.lib.push_notifications import uses_notification_bouncer + + if uses_notification_bouncer(): + # Let the bouncer know about the new realm. + # We do this in a queue worker to avoid messing with the realm + # creation process due to network issues or latency. + event = {"type": "register_realm_with_push_bouncer", "realm_id": realm.id} + queue_json_publish("deferred_work", event) diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py index dc531d31f0..5b7c683bc7 100644 --- a/zerver/tests/test_import_export.py +++ b/zerver/tests/test_import_export.py @@ -1,4 +1,5 @@ import datetime +import json import os import shutil import uuid @@ -10,6 +11,7 @@ import orjson from django.conf import settings from django.core.exceptions import ValidationError from django.db.models import Q, QuerySet +from django.test import override_settings from django.utils.timezone import now as timezone_now from typing_extensions import override @@ -1388,6 +1390,25 @@ class RealmImportExportTest(ExportFile): self.assertEqual(realm_user_default.default_language, "en") self.assertEqual(realm_user_default.twenty_four_hour_time, False) + @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com") + def test_import_realm_notify_bouncer(self) -> None: + original_realm = Realm.objects.get(string_id="zulip") + + self.export_realm(original_realm) + + with self.settings(BILLING_ENABLED=False), self.assertLogs(level="INFO"), patch( + "zerver.lib.remote_server.send_to_push_bouncer" + ) as m: + new_realm = do_import_realm(get_output_dir(), "test-zulip") + + self.assertTrue(Realm.objects.filter(string_id="test-zulip").exists()) + calls_args_for_assert = m.call_args_list[0][0] + self.assertEqual(calls_args_for_assert[0], "POST") + self.assertEqual(calls_args_for_assert[1], "server/analytics") + self.assertIn( + new_realm.id, [realm["id"] for realm in json.loads(m.call_args_list[0][0][2]["realms"])] + ) + def test_import_files_from_local(self) -> None: user = self.example_user("hamlet") realm = user.realm diff --git a/zerver/tests/test_realm.py b/zerver/tests/test_realm.py index bc7f5108e3..80b242023e 100644 --- a/zerver/tests/test_realm.py +++ b/zerver/tests/test_realm.py @@ -1,4 +1,5 @@ import datetime +import json import os import re from datetime import timedelta @@ -7,6 +8,7 @@ from unittest import mock import orjson from django.conf import settings +from django.test import override_settings from django.utils.timezone import now as timezone_now from typing_extensions import override @@ -1061,6 +1063,21 @@ class RealmTest(ZulipTestCase): ] self.assertEqual(sorted(user_group_names), sorted(expected_system_group_names)) + @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com") + def test_do_create_realm_notify_bouncer(self) -> None: + with mock.patch("zerver.lib.remote_server.send_to_push_bouncer") as m: + realm = do_create_realm("realm_string_id", "realm name") + + self.assertEqual(realm.string_id, "realm_string_id") + self.assertEqual(m.call_count, 1) + + calls_args_for_assert = m.call_args_list[0][0] + self.assertEqual(calls_args_for_assert[0], "POST") + self.assertEqual(calls_args_for_assert[1], "server/analytics") + self.assertIn( + realm.id, [realm["id"] for realm in json.loads(m.call_args_list[0][0][2]["realms"])] + ) + def test_changing_waiting_period_updates_system_groups(self) -> None: realm = get_realm("zulip") members_system_group = UserGroup.objects.get( diff --git a/zerver/worker/queue_processors.py b/zerver/worker/queue_processors.py index 63ef4d9a65..0137254386 100644 --- a/zerver/worker/queue_processors.py +++ b/zerver/worker/queue_processors.py @@ -78,7 +78,10 @@ from zerver.lib.push_notifications import ( ) from zerver.lib.pysa import mark_sanitized from zerver.lib.queue import SimpleQueueClient, retry_event -from zerver.lib.remote_server import PushNotificationBouncerRetryLaterError +from zerver.lib.remote_server import ( + PushNotificationBouncerRetryLaterError, + send_realms_only_to_push_bouncer, +) from zerver.lib.send_email import ( EmailNotDeliveredError, FromAddress, @@ -1167,6 +1170,13 @@ class DeferredWorker(QueueProcessingWorker): ) user_profile = get_user_profile_by_id(event["user_profile_id"]) reactivate_user_if_soft_deactivated(user_profile) + elif event["type"] == "register_realm_with_push_bouncer": + # In the future we may use the realm_id to send only that single realm's info. + realm_id = event["realm_id"] + logger.info( + "Running send_realms_only_to_push_bouncer, requested due to realm %s", realm_id + ) + send_realms_only_to_push_bouncer() end = time.time() logger.info(