mirror of https://github.com/zulip/zulip.git
86 lines
3.5 KiB
Python
86 lines
3.5 KiB
Python
from argparse import ArgumentParser
|
|
from typing import Any, List
|
|
|
|
from zerver.actions.streams import (
|
|
bulk_add_subscriptions,
|
|
bulk_remove_subscriptions,
|
|
do_deactivate_stream,
|
|
)
|
|
from zerver.lib.cache import cache_delete_many, to_dict_cache_key_id
|
|
from zerver.lib.management import ZulipBaseCommand
|
|
from zerver.models import Message, Subscription, get_stream
|
|
|
|
|
|
def bulk_delete_cache_keys(message_ids_to_clear: List[int]) -> None:
|
|
while len(message_ids_to_clear) > 0:
|
|
batch = message_ids_to_clear[0:5000]
|
|
|
|
keys_to_delete = [to_dict_cache_key_id(message_id) for message_id in batch]
|
|
cache_delete_many(keys_to_delete)
|
|
|
|
message_ids_to_clear = message_ids_to_clear[5000:]
|
|
|
|
|
|
class Command(ZulipBaseCommand):
|
|
help = """Merge two streams."""
|
|
|
|
def add_arguments(self, parser: ArgumentParser) -> None:
|
|
parser.add_argument("stream_to_keep", help="name of stream to keep")
|
|
parser.add_argument(
|
|
"stream_to_destroy", help="name of stream to merge into the stream being kept"
|
|
)
|
|
self.add_realm_args(parser, required=True)
|
|
|
|
def handle(self, *args: Any, **options: str) -> None:
|
|
realm = self.get_realm(options)
|
|
assert realm is not None # Should be ensured by parser
|
|
stream_to_keep = get_stream(options["stream_to_keep"], realm)
|
|
stream_to_destroy = get_stream(options["stream_to_destroy"], realm)
|
|
|
|
recipient_to_destroy = stream_to_destroy.recipient
|
|
recipient_to_keep = stream_to_keep.recipient
|
|
|
|
# The high-level approach here is to move all the messages to
|
|
# the surviving stream, deactivate all the subscriptions on
|
|
# the stream to be removed and deactivate the stream, and add
|
|
# new subscriptions to the stream to keep for any users who
|
|
# were only on the now-deactivated stream.
|
|
|
|
# Move the messages, and delete the old copies from caches.
|
|
message_ids_to_clear = list(
|
|
Message.objects.filter(recipient=recipient_to_destroy).values_list("id", flat=True)
|
|
)
|
|
count = Message.objects.filter(recipient=recipient_to_destroy).update(
|
|
recipient=recipient_to_keep
|
|
)
|
|
print(f"Moved {count} messages")
|
|
bulk_delete_cache_keys(message_ids_to_clear)
|
|
|
|
# Move the Subscription objects. This algorithm doesn't
|
|
# preserve any stream settings/colors/etc. from the stream
|
|
# being destroyed, but it's convenient.
|
|
existing_subs = Subscription.objects.filter(recipient=recipient_to_keep)
|
|
users_already_subscribed = {sub.user_profile_id: sub.active for sub in existing_subs}
|
|
|
|
subs_to_deactivate = Subscription.objects.filter(
|
|
recipient=recipient_to_destroy, active=True
|
|
)
|
|
users_to_activate = [
|
|
sub.user_profile
|
|
for sub in subs_to_deactivate
|
|
if not users_already_subscribed.get(sub.user_profile_id, False)
|
|
]
|
|
|
|
if len(subs_to_deactivate) > 0:
|
|
print(f"Deactivating {len(subs_to_deactivate)} subscriptions")
|
|
bulk_remove_subscriptions(
|
|
realm,
|
|
[sub.user_profile for sub in subs_to_deactivate],
|
|
[stream_to_destroy],
|
|
acting_user=None,
|
|
)
|
|
do_deactivate_stream(stream_to_destroy, acting_user=None)
|
|
if len(users_to_activate) > 0:
|
|
print(f"Adding {len(users_to_activate)} subscriptions")
|
|
bulk_add_subscriptions(realm, [stream_to_keep], users_to_activate, acting_user=None)
|