From 5a32ea52aeaeae7e98974c50d91a74f4405a107a Mon Sep 17 00:00:00 2001 From: Alex Vandiver Date: Thu, 3 Aug 2023 20:20:37 +0000 Subject: [PATCH] send_custom_email: Stop turning every user query into an id-based set. The set of objects in the `users` object can be very large (in some cases, literally every object in the database) and making them into a giant `id in (...)` to handle the one tiny corner case which we never use is silly. Switch the `--users` codepath to returning a QuerySet as well, so it can be composed. We pass a QuerySet into send_custom_email as well, so it can ensure that the realm is `select_related` in as well, no matter how the QuerySet was generated. --- zerver/lib/management.py | 24 +++++++++++++++--- zerver/lib/send_email.py | 7 +++--- .../management/commands/send_custom_email.py | 17 ++++--------- .../commands/sync_ldap_user_data.py | 5 ++-- zerver/tests/test_email_notifications.py | 25 ++++++++----------- zerver/tests/test_management_commands.py | 24 ++++++++++++++++-- 6 files changed, 63 insertions(+), 39 deletions(-) diff --git a/zerver/lib/management.py b/zerver/lib/management.py index 8e7c509aa6..36beef5c03 100644 --- a/zerver/lib/management.py +++ b/zerver/lib/management.py @@ -2,12 +2,14 @@ import logging from argparse import ArgumentParser, RawTextHelpFormatter, _ActionsContainer from dataclasses import dataclass -from typing import Any, Collection, Dict, Optional +from functools import reduce +from typing import Any, Dict, Optional from django.conf import settings from django.core import validators from django.core.exceptions import MultipleObjectsReturned, ValidationError from django.core.management.base import BaseCommand, CommandError, CommandParser +from django.db.models import Q, QuerySet from zerver.lib.initial_password import initial_password from zerver.models import Client, Realm, UserProfile, get_client @@ -115,7 +117,7 @@ server via `ps -ef` or reading bash history. Prefer realm: Optional[Realm], is_bot: Optional[bool] = None, include_deactivated: bool = False, - ) -> Collection[UserProfile]: + ) -> QuerySet[UserProfile]: if "all_users" in options: all_users = options["all_users"] @@ -137,9 +139,23 @@ server via `ps -ef` or reading bash history. Prefer return user_profiles if options["users"] is None: - return [] + return UserProfile.objects.none() emails = {email.strip() for email in options["users"].split(",")} - return [self.get_user(email, realm) for email in emails] + # This is inefficient, but we fetch (and throw away) the + # get_user of each email, so that we verify that the email + # address/realm returned only one result; it may return more + # if realm is not specified but email address was. + for email in emails: + self.get_user(email, realm) + + user_profiles = UserProfile.objects.all().select_related("realm") + if realm is not None: + user_profiles = user_profiles.filter(realm=realm) + email_matches = [Q(delivery_email__iexact=e) for e in emails] + user_profiles = user_profiles.filter(reduce(lambda a, b: a | b, email_matches)) + + # Return the single query, for ease of composing. + return user_profiles def get_user(self, email: str, realm: Optional[Realm]) -> UserProfile: # If a realm is specified, try to find the user there, and diff --git a/zerver/lib/send_email.py b/zerver/lib/send_email.py index ff34e8afc7..8224a5fbf3 100644 --- a/zerver/lib/send_email.py +++ b/zerver/lib/send_email.py @@ -8,7 +8,7 @@ from email.headerregistry import Address from email.parser import Parser from email.policy import default from email.utils import formataddr, parseaddr -from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union import backoff import css_inline @@ -20,6 +20,7 @@ from django.core.mail.backends.smtp import EmailBackend from django.core.mail.message import sanitize_address from django.core.management import CommandError from django.db import transaction +from django.db.models import QuerySet from django.http import HttpRequest from django.template import loader from django.utils.timezone import now as timezone_now @@ -505,7 +506,7 @@ def get_header(option: Optional[str], header: Optional[str], name: str) -> str: def send_custom_email( - users: Iterable[UserProfile], *, target_emails: Sequence[str] = [], options: Dict[str, Any] + users: QuerySet[UserProfile], *, target_emails: Sequence[str] = [], options: Dict[str, Any] ) -> None: """ Helper for `manage.py send_custom_email`. @@ -554,7 +555,7 @@ def send_custom_email( f.write(get_header(options.get("subject"), parsed_email_template.get("subject"), "subject")) # Finally, we send the actual emails. - for user_profile in users: + for user_profile in users.select_related("realm"): if options.get("admins_only") and not user_profile.is_realm_admin: continue context = { diff --git a/zerver/management/commands/send_custom_email.py b/zerver/management/commands/send_custom_email.py index 42c885d9cb..b7a27d7e3f 100644 --- a/zerver/management/commands/send_custom_email.py +++ b/zerver/management/commands/send_custom_email.py @@ -1,9 +1,8 @@ from argparse import ArgumentParser -from typing import Any, Collection, List +from typing import Any, List from django.conf import settings -from django.core.management.base import CommandError -from django.db.models import Q +from django.db.models import Q, QuerySet from zerver.lib.management import ZulipBaseCommand from zerver.lib.send_email import send_custom_email @@ -79,7 +78,7 @@ class Command(ZulipBaseCommand): def handle(self, *args: Any, **options: str) -> None: target_emails: List[str] = [] - users: Collection[UserProfile] = [] + users: QuerySet[UserProfile] = UserProfile.objects.none() if options["entire_server"]: users = UserProfile.objects.filter( @@ -127,14 +126,8 @@ class Command(ZulipBaseCommand): # Only email users who've agreed to the terms of service. if settings.TERMS_OF_SERVICE_VERSION is not None: - # We need to do a new query because the `get_users` path - # passes us a list rather than a QuerySet. - users = ( - UserProfile.objects.select_related("realm") - .filter(id__in=[u.id for u in users]) - .exclude( - Q(tos_version=None) | Q(tos_version=UserProfile.TOS_VERSION_BEFORE_FIRST_LOGIN) - ) + users = users.exclude( + Q(tos_version=None) | Q(tos_version=UserProfile.TOS_VERSION_BEFORE_FIRST_LOGIN) ) send_custom_email(users, target_emails=target_emails, options=options) diff --git a/zerver/management/commands/sync_ldap_user_data.py b/zerver/management/commands/sync_ldap_user_data.py index a8a3bd259a..8b9d7a7770 100644 --- a/zerver/management/commands/sync_ldap_user_data.py +++ b/zerver/management/commands/sync_ldap_user_data.py @@ -78,13 +78,12 @@ class Command(ZulipBaseCommand): realm = self.get_realm(options) user_profiles = self.get_users(options, realm, is_bot=False, include_deactivated=True) else: - user_profile_query = UserProfile.objects.select_related("realm").filter(is_bot=False) + user_profiles = UserProfile.objects.select_related("realm").filter(is_bot=False) - if not user_profile_query.exists(): + if not user_profiles.exists(): # This case provides a special error message if one # tries setting up LDAP sync before creating a realm. raise CommandError("Zulip server contains no users. Have you created a realm?") - user_profiles = list(user_profile_query) if len(user_profiles) == 0: # We emphasize that this error is purely about the diff --git a/zerver/tests/test_email_notifications.py b/zerver/tests/test_email_notifications.py index 8a15fb0cb3..b757a23869 100644 --- a/zerver/tests/test_email_notifications.py +++ b/zerver/tests/test_email_notifications.py @@ -18,12 +18,7 @@ from zerver.lib.email_notifications import ( ) from zerver.lib.send_email import deliver_scheduled_emails, send_custom_email from zerver.lib.test_classes import ZulipTestCase -from zerver.models import ( - Realm, - ScheduledEmail, - UserProfile, - get_realm, -) +from zerver.models import Realm, ScheduledEmail, UserProfile, get_realm class TestCustomEmails(ZulipTestCase): @@ -37,7 +32,7 @@ class TestCustomEmails(ZulipTestCase): markdown_template.write(b"# Some heading\n\nSome content\n{{ realm_name }}") markdown_template.flush() send_custom_email( - [hamlet], + UserProfile.objects.filter(id=hamlet.id), options={ "markdown_template_path": markdown_template.name, "reply_to": reply_to, @@ -68,7 +63,7 @@ class TestCustomEmails(ZulipTestCase): contact_email = "zulip-admin@example.com" markdown_template_path = "templates/corporate/policies/index.md" send_custom_email( - [], + UserProfile.objects.none(), target_emails=[contact_email], options={ "markdown_template_path": markdown_template_path, @@ -98,7 +93,7 @@ class TestCustomEmails(ZulipTestCase): "zerver/tests/fixtures/email/custom_emails/email_base_headers_test.html" ) send_custom_email( - [hamlet], + UserProfile.objects.filter(id=hamlet.id), options={ "markdown_template_path": markdown_template_path, "dry_run": False, @@ -123,7 +118,7 @@ class TestCustomEmails(ZulipTestCase): self.assertRaises( NoEmailArgumentError, send_custom_email, - [hamlet], + UserProfile.objects.filter(id=hamlet.id), options={ "markdown_template_path": markdown_template_path, "from_name": from_name, @@ -134,7 +129,7 @@ class TestCustomEmails(ZulipTestCase): self.assertRaises( NoEmailArgumentError, send_custom_email, - [hamlet], + UserProfile.objects.filter(id=hamlet.id), options={ "markdown_template_path": markdown_template_path, "subject": email_subject, @@ -155,7 +150,7 @@ class TestCustomEmails(ZulipTestCase): self.assertRaises( DoubledEmailArgumentError, send_custom_email, - [hamlet], + UserProfile.objects.filter(id=hamlet.id), options={ "markdown_template_path": markdown_template_path, "subject": email_subject, @@ -166,7 +161,7 @@ class TestCustomEmails(ZulipTestCase): self.assertRaises( DoubledEmailArgumentError, send_custom_email, - [hamlet], + UserProfile.objects.filter(id=hamlet.id), options={ "markdown_template_path": markdown_template_path, "from_name": from_name, @@ -184,7 +179,7 @@ class TestCustomEmails(ZulipTestCase): "zerver/tests/fixtures/email/custom_emails/email_base_headers_test.html" ) send_custom_email( - [admin_user, non_admin_user], + UserProfile.objects.filter(id__in=(admin_user.id, non_admin_user.id)), options={ "markdown_template_path": markdown_template_path, "admins_only": True, @@ -202,7 +197,7 @@ class TestCustomEmails(ZulipTestCase): markdown_template_path = "templates/zerver/tests/markdown/test_nested_code_blocks.md" with patch("builtins.print") as _: send_custom_email( - [hamlet], + UserProfile.objects.filter(id=hamlet.id), options={ "markdown_template_path": markdown_template_path, "reply_to": reply_to, diff --git a/zerver/tests/test_management_commands.py b/zerver/tests/test_management_commands.py index 72db9e9b76..8000d409ce 100644 --- a/zerver/tests/test_management_commands.py +++ b/zerver/tests/test_management_commands.py @@ -129,11 +129,11 @@ class TestZulipBaseCommand(ZulipTestCase): self.command.get_users(dict(users=user_emails), self.zulip_realm) self.assertEqual( - self.command.get_users(dict(users=self.example_email("iago")), self.zulip_realm), + list(self.command.get_users(dict(users=self.example_email("iago")), self.zulip_realm)), [self.example_user("iago")], ) - self.assertEqual(self.command.get_users(dict(users=None), None), []) + self.assertEqual(list(self.command.get_users(dict(users=None), None)), []) def test_get_users_with_all_users_argument_enabled(self) -> None: expected_user_profiles = self.sorted_users( @@ -607,6 +607,7 @@ class TestSendCustomEmail(ZulipTestCase): def test_custom_email_with_dry_run(self) -> None: path = "templates/zerver/tests/markdown/test_nested_code_blocks.md" user = self.example_user("hamlet") + other_user = self.example_user("cordelia") with patch("builtins.print") as mock_print: call_command( @@ -625,3 +626,22 @@ class TestSendCustomEmail(ZulipTestCase): call(" hamlet@zulip.com (zulip)"), ], ) + + with patch("builtins.print") as mock_print: + call_command( + self.COMMAND_NAME, + "-r=zulip", + f"--path={path}", + f"-u={user.delivery_email},{other_user.delivery_email}", + "--subject=Test email", + "--from-name=zulip@zulip.example.com", + "--dry-run", + ) + self.assertEqual( + mock_print.mock_calls[1:], + [ + call("Would send the above email to:"), + call(" cordelia@zulip.com (zulip)"), + call(" hamlet@zulip.com (zulip)"), + ], + )