From c80802ff1e91b7ed1db55294b24e83b1c9346faa Mon Sep 17 00:00:00 2001 From: Greg Price Date: Tue, 6 Feb 2018 18:59:57 -0800 Subject: [PATCH] bots: Clean up create_realm_internal_bots. This code duplicated the code in setup_realm_internal_bots, with some added logic to avoid trying to create the same bot twice. That logic was buggy so that it would never work at all -- it subtracted a set of UserProfile objects from a set of email strings -- so it looked like the command might blow up when run after the users already existed. In fact, the buggy logic wasn't necessary, because the work the command does after it is idempotent -- in particular `create_users`, within its subroutine `bulk_create_users`, already filters out users that already exist. So just cut the buggy stuff out, deduplicate the rest with `setup_realm_internal_bots`, and document that invariant on the latter. While we're here, in the common case bail early without doing any per-realm work in Python, since we're running this on every upgrade. --- zerver/lib/onboarding.py | 5 ++ .../commands/create_realm_internal_bots.py | 52 ++++++++----------- 2 files changed, 28 insertions(+), 29 deletions(-) diff --git a/zerver/lib/onboarding.py b/zerver/lib/onboarding.py index d380d8d7bd..bb0869bdb8 100644 --- a/zerver/lib/onboarding.py +++ b/zerver/lib/onboarding.py @@ -10,6 +10,11 @@ from zerver.models import Realm, UserProfile, Message, Reaction, get_system_bot from typing import Any, Dict, List, Mapping, Text def setup_realm_internal_bots(realm: Realm) -> None: + """Create this realm's internal bots. + + This function is idempotent; it does nothing for a bot that + already exists. + """ internal_bots = [(bot['name'], bot['email_template'] % (settings.INTERNAL_BOT_DOMAIN,)) for bot in settings.REALM_INTERNAL_BOTS] create_users(realm, internal_bots, bot_type=UserProfile.DEFAULT_BOT) diff --git a/zerver/management/commands/create_realm_internal_bots.py b/zerver/management/commands/create_realm_internal_bots.py index b0e1dc55ae..a83f7f41b8 100644 --- a/zerver/management/commands/create_realm_internal_bots.py +++ b/zerver/management/commands/create_realm_internal_bots.py @@ -3,38 +3,32 @@ from typing import Any, Iterable, Text, Tuple from django.conf import settings from django.core.management.base import BaseCommand +from django.db.models import Count -from zerver.lib.actions import create_users +from zerver.lib.onboarding import setup_realm_internal_bots from zerver.models import Realm, UserProfile class Command(BaseCommand): - help = "Create Realm internal bots. These bots provide various services like doing reminders." + help = """\ +Create realm internal bots if absent, in all realms. + +These are normally created when the realm is, so this should be a no-op +except when upgrading to a version that adds a new realm internal bot. +""" + + @staticmethod + def missing_any_bots() -> bool: + bot_emails = [bot['email_template'] % (settings.INTERNAL_BOT_DOMAIN,) + for bot in settings.REALM_INTERNAL_BOTS] + bot_counts = dict(UserProfile.objects.filter(email__in=bot_emails) + .values_list('email') + .annotate(Count('id'))) + realm_count = Realm.objects.count() + return any(bot_counts.get(email, 0) < realm_count for email in bot_emails) def handle(self, *args: Any, **options: Any) -> None: - internal_bots = set([(bot['name'], bot['email_template'] % (settings.INTERNAL_BOT_DOMAIN,)) - for bot in settings.REALM_INTERNAL_BOTS]) - - existing_bots = list(UserProfile.objects.select_related( - 'realm').filter(email__in=[bot[1] for bot in internal_bots])) - - all_realms = list(Realm.objects.all()) - - for realm in all_realms: - this_realm_bots = set() - for bot in existing_bots: - if bot.realm.string_id == realm.string_id: - this_realm_bots.update([bot]) - bots_to_create = list(internal_bots - this_realm_bots) - if bots_to_create: - create_users(realm, bots_to_create, bot_type=UserProfile.DEFAULT_BOT) - - # Set the owners for these bots to the bots themselves - bots = UserProfile.objects.filter( - email__in=[bot_info[1] for bot_info in internal_bots], - bot_owner__isnull=True - ) - for bot in bots: - bot.bot_owner = bot - bot.save() - - self.stdout.write("Successfully created realm default bots.\n") + if self.missing_any_bots(): + for realm in Realm.objects.all(): + setup_realm_internal_bots(realm) + # create_users is idempotent -- it's a no-op when a given email + # already has a user in a given realm.