zulip/zerver/management/commands/restore_messages.py

from typing import Any

from django.core.management.base import CommandParser

from zerver.lib.management import ZulipBaseCommand
from zerver.lib.retention import (
    restore_all_data_from_archive,
    restore_data_from_archive,
    restore_data_from_archive_by_realm,
)
from zerver.models import ArchiveTransaction


class Command(ZulipBaseCommand):
    help = """
Restore recently deleted messages from the archive, that
have not been vacuumed (because the time limit of
ARCHIVED_DATA_VACUUMING_DELAY_DAYS has not passed).

Intended primarily for use after against potential bugs in
Zulip's message retention and deletion features.

Examples:
To restore all recently deleted messages:
  ./manage.py restore_messages
To restore a specific ArchiveTransaction:
  ./manage.py restore_messages --transaction-id=1
"""

    def add_arguments(self, parser: CommandParser) -> None:
        parser.add_argument(
            "-d",
            "--restore-deleted",
            action="store_true",
            help="Restore manually deleted messages.",
        )
        parser.add_argument(
            "-t", "--transaction-id", type=int, help="Restore a specific ArchiveTransaction."
        )

        self.add_realm_args(
            parser,
            help="Restore archived messages from the specified realm. "
            "(Does not restore manually deleted messages.)",
        )

    def handle(self, **options: Any) -> None:
        realm = self.get_realm(options)
        if realm:
            restore_data_from_archive_by_realm(realm)
        elif options["transaction_id"]:
            restore_data_from_archive(ArchiveTransaction.objects.get(id=options["transaction_id"]))
        else:
            restore_all_data_from_archive(restore_manual_transactions=options["restore_deleted"])