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.