zulip/zerver/management/commands/fix_unreads.py

61 lines
2.0 KiB
Python

import logging
from argparse import ArgumentParser
from typing import Any, List, Optional
from django.db import connection
from zerver.lib.fix_unreads import fix
from zerver.lib.management import CommandError, ZulipBaseCommand
from zerver.models import Realm, UserProfile
logging.getLogger('zulip.fix_unreads').setLevel(logging.INFO)
class Command(ZulipBaseCommand):
help = """Fix problems related to unread counts."""
def add_arguments(self, parser: ArgumentParser) -> None:
parser.add_argument('emails',
metavar='<emails>',
type=str,
nargs='*',
help='email address to spelunk')
parser.add_argument('--all',
action='store_true',
dest='all',
default=False,
help='fix all users in specified realm')
self.add_realm_args(parser)
def fix_all_users(self, realm: Realm) -> None:
user_profiles = list(UserProfile.objects.filter(
realm=realm,
is_bot=False
))
for user_profile in user_profiles:
fix(user_profile)
connection.commit()
def fix_emails(self, realm: Optional[Realm], emails: List[str]) -> None:
for email in emails:
try:
user_profile = self.get_user(email, realm)
except CommandError:
print(f"e-mail {email} doesn't exist in the realm {realm}, skipping")
return
fix(user_profile)
connection.commit()
def handle(self, *args: Any, **options: Any) -> None:
realm = self.get_realm(options)
if options['all']:
if realm is None:
raise CommandError('You must specify a realm if you choose the --all option.')
self.fix_all_users(realm)
return
self.fix_emails(realm, options['emails'])