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

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

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 = {

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, QuerySet
-from django.db.models import Q
 
 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
+            users = users.exclude(
-            # passes us a list rather than a QuerySet.
+                Q(tos_version=None) | Q(tos_version=UserProfile.TOS_VERSION_BEFORE_FIRST_LOGIN)
-            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)
-                )
             )
         send_custom_email(users, target_emails=target_emails, options=options)
 

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

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 (
+from zerver.models import Realm, ScheduledEmail, UserProfile, get_realm
-    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,

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)"),
+                ],
+            )