zulip/zerver/management/commands/export_usermessage_batch.py

49 lines
1.8 KiB
Python
Raw Normal View History

from typing import Any
from argparse import ArgumentParser
from django.core.management.base import BaseCommand
import glob
import logging
import os
import shutil
from zerver.lib.export import export_usermessages_batch
class Command(BaseCommand):
help = """UserMessage fetching helper for export.py"""
def add_arguments(self, parser):
# type: (ArgumentParser) -> None
parser.add_argument('--path',
dest='path',
action="store",
default=None,
help='Path to find messages.json archives')
parser.add_argument('--thread',
dest='thread',
action="store",
default=None,
help='Thread ID')
def handle(self, *args, **options):
# type: (*Any, **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)
except Exception:
# Put the item back in the free pool when we fail
shutil.move(locked_path, partial_path)
raise