Mirror
/
zulip
mirror of https://github.com/zulip/zulip.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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.
This commit is contained in:
Alex Vandiver 2023-08-03 20:20:37 +00:00 committed by Tim Abbott
parent 6b25fab38c
commit 5a32ea52ae
6 changed files with 63 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
zerver/lib/management.py
View File

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

7
zerver/lib/send_email.py
View File

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

17
zerver/management/commands/send_custom_email.py
View File

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

5
zerver/management/commands/sync_ldap_user_data.py
View File

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

25
zerver/tests/test_email_notifications.py
View File

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

24
zerver/tests/test_management_commands.py
View File

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