zulip/zerver/management/commands/export_usermessage_batch.py

49 lines
1.9 KiB
Python

import glob
import logging
import os
import shutil
from argparse import ArgumentParser
from typing import Any
from django.core.management.base import BaseCommand
from zerver.lib.export import export_usermessages_batch
class Command(BaseCommand):
help = """UserMessage fetching helper for export.py"""
def add_arguments(self, parser: ArgumentParser) -> None:
parser.add_argument('--path',
action="store",
default=None,
help='Path to find messages.json archives')
parser.add_argument('--thread',
action="store",
default=None,
help='Thread ID')
parser.add_argument('--consent-message-id',
action="store",
default=None,
type=int,
help='ID of the message advertising users to react with thumbs up')
def handle(self, *args: Any, **options: Any) -> None:
logging.info("Starting UserMessage batch thread %s", options['thread'])
files = set(glob.glob(os.path.join(options['path'], 'messages-*.json.partial')))
for partial_path in files:
locked_path = partial_path.replace(".json.partial", ".json.locked")
output_path = partial_path.replace(".json.partial", ".json")
try:
shutil.move(partial_path, locked_path)
except Exception:
# Already claimed by another process
continue
logging.info("Thread %s processing %s", options['thread'], output_path)
try:
export_usermessages_batch(locked_path, output_path, options["consent_message_id"])
except Exception:
# Put the item back in the free pool when we fail
shutil.move(locked_path, partial_path)
raise