zulip/zerver/lib/management.py

# Library code for use in management commands
from __future__ import absolute_import
from __future__ import print_function

from argparse import ArgumentParser
from django.core.exceptions import MultipleObjectsReturned
from django.core.management.base import BaseCommand, CommandError
from typing import Any, Dict, Optional, Text

from zerver.models import Realm, UserProfile

def is_integer_string(val):
    # type: (str) -> bool
    try:
        int(val)
        return True
    except ValueError:
        return False

class ZulipBaseCommand(BaseCommand):
    def add_realm_args(self, parser, required=False, help=None):
        # type: (ArgumentParser, bool, Optional[str]) -> None
        if help is None:
            help = """The numeric or string ID (subdomain) of the Zulip organization to modify.
You can use the command list_realms to find ID of the realms in this server."""

        parser.add_argument(
            '-r', '--realm',
            dest='realm_id',
            required=required,
            type=str,
            help=help)

    def get_realm(self, options):
        # type: (Dict[str, Any]) -> Optional[Realm]
        val = options["realm_id"]
        if val is None:
            return None

        # If they specified a realm argument, we need to ensure the
        # realm exists.  We allow two formats: the numeric ID for the
        # realm and the string ID of the realm.
        try:
            if is_integer_string(val):
                return Realm.objects.get(id=val)
            return Realm.objects.get(string_id=val)
        except Realm.DoesNotExist:
            raise CommandError("There is no realm with id '%s'. Aborting." %
                               (options["realm_id"],))

    def get_user(self, email, realm):
        # type: (Text, Optional[Realm]) -> UserProfile

        # If a realm is specified, try to find the user there, and
        # throw an error if they don't exist.
        if realm is not None:
            try:
                return UserProfile.objects.select_related().get(email__iexact=email.strip(), realm=realm)
            except UserProfile.DoesNotExist:
                raise CommandError("The realm '%s' does not contain a user with email '%s'" % (realm, email))

        # Realm is None in the remaining code path.  Here, we
        # optimistically try to see if there is exactly one user with
        # that email; if so, we'll return it.
        try:
            return UserProfile.objects.select_related().get(email__iexact=email.strip())
        except MultipleObjectsReturned:
            raise CommandError("This Zulip server contains multiple users with that email " +
                               "(in different realms); please pass `--realm` to specify which one to modify.")
        except UserProfile.DoesNotExist:
            raise CommandError("This Zulip server does not contain a user with email '%s'" % (email,))