<commit_before>import datetime
import gzip
import inspect
import json
import logging
import os
import pathlib
import re
import shutil
import zipfile
from collections import namedtuple
from concurrent import futures
from pathlib import Path

import atexit
from django.core.management import BaseCommand, CommandError

from corehq.apps.dump_reload.management.commands.load_domain_data import get_tmp_extract_dir
from corehq.blobs.export import BlobDbBackendExporter
from corehq.blobs.management.commands.run_blob_export import get_lines_from_file
from corehq.util.decorators import change_log_level
from corehq.util.log import with_progress_bar

BlobMetaKey = namedtuple('BlobMetaKey', 'key')


class Command(BaseCommand):
    help = inspect.cleandoc("""
    Usage ./manage.py export_selected_blobs [options] path_to_export_zip

    'path_to_export_zip' must be a ZIP file generated using `dump_domain_data` and must
    include a `meta.json` file with the object counts.

    To top-up an older blob dump, first extract a list of names from the archive:
        $ tar --list -f blob_export.tar.gz > blob_export.list
    Then provide this file to the `--already_exported` argument to skip over
    those objects in this dump.
    """)

    def add_arguments(self, parser):
        parser.add_argument('path', help='Path to a file with a list of blob meta JSON to export')
        parser.add_argument('--meta-file-filter', nargs='?',
                            help='Only export blobs for metadata in files with filenames matching the filter regex.')
        parser.add_argument('--already_exported', dest='already_exported',
                            help='Pass a file with a list of blob names already exported')
        parser.add_argument('--json-output', action="store_true", help="Produce JSON output for use in tests")

    @change_log_level('boto3', logging.WARNING)
    @change_log_level('botocore', logging.WARNING)
    def handle(self, path, **options):
        already_exported = get_lines_from_file(options['already_exported'])
        if already_exported:
            print("Found {} existing blobs, these will be skipped".format(len(already_exported)))

        filter_pattern = options.get('meta-file-filter')
        filter_rx = None
        if filter_pattern:
            filter_rx = re.compile(filter_pattern)

        def _filter(filename):
            return 'blob_meta' in filename and (not filter_rx or filter_rx.match(filename))

        target_dir = get_tmp_extract_dir(path, specifier='blob_meta')
        atexit.register(lambda: shutil.rmtree(target_dir))

        target_path = Path(target_dir)
        export_meta_files = []
        with zipfile.ZipFile(path, 'r') as archive:
            meta = json.loads(archive.read("meta.json"))
            for dump_file in archive.namelist():
                if _filter(dump_file):
                    export_meta_files.append(target_path.joinpath(dump_file))
                    if not target_path.joinpath(dump_file).exists():
                        archive.extract(dump_file, target_dir)

        results = []
        filenames = []
        with futures.ThreadPoolExecutor(max_workers=len(export_meta_files)) as executor:
            for path in export_meta_files:
                results.append(executor.submit(_export_blobs, path, meta))

            for result in futures.as_completed(results):<commit_msg>option to use extracted archives<commit_after>import datetime
import gzip
import inspect
import json
import logging
import os
import pathlib
import re
import shutil
import zipfile
from collections import namedtuple
from concurrent import futures
from pathlib import Path

import atexit
from django.core.management import BaseCommand, CommandError

from corehq.apps.dump_reload.management.commands.load_domain_data import get_tmp_extract_dir
from corehq.blobs.export import BlobDbBackendExporter
from corehq.blobs.management.commands.run_blob_export import get_lines_from_file
from corehq.util.decorators import change_log_level
from corehq.util.log import with_progress_bar

BlobMetaKey = namedtuple('BlobMetaKey', 'key')


class Command(BaseCommand):
    help = inspect.cleandoc("""
    Usage ./manage.py export_selected_blobs [options] path_to_export_zip

    'path_to_export_zip' must be a ZIP file generated using `dump_domain_data` and must
    include a `meta.json` file with the object counts.

    To top-up an older blob dump, first extract a list of names from the archive:
        $ tar --list -f blob_export.tar.gz > blob_export.list
    Then provide this file to the `--already_exported` argument to skip over
    those objects in this dump.
    """)

    def add_arguments(self, parser):
        parser.add_argument('path', help='Path to a file with a list of blob meta JSON to export')
        parser.add_argument('--meta-file-filter', nargs='?',
                            help='Only export blobs for metadata in files with filenames matching the filter regex.')
        parser.add_argument('--already_exported', dest='already_exported',
                            help='Pass a file with a list of blob names already exported')
        parser.add_argument('--use-extracted', action='store_true', default=False, dest='use_extracted',
                            help="Use already extracted dump if it exists.")
        parser.add_argument('--json-output', action="store_true", help="Produce JSON output for use in tests")

    @change_log_level('boto3', logging.WARNING)
    @change_log_level('botocore', logging.WARNING)
    def handle(self, path, **options):
        self.use_extracted = options.get('use_extracted')

        already_exported = get_lines_from_file(options['already_exported'])
        if already_exported:
            print("Found {} existing blobs, these will be skipped".format(len(already_exported)))

        filter_pattern = options.get('meta-file-filter')
        filter_rx = None
        if filter_pattern:
            filter_rx = re.compile(filter_pattern)

        def _filter(filename):
            return 'blob_meta' in filename and (not filter_rx or filter_rx.match(filename))

        target_dir = get_tmp_extract_dir(path, specifier='blob_meta')
        target_path = Path(target_dir)
        export_meta_files = []
        if not target_path.exists():
            with zipfile.ZipFile(path, 'r') as archive:
                meta = json.loads(archive.read("meta.json"))
                for dump_file in archive.namelist():
                    if _filter(dump_file):
                        export_meta_files.append(target_path.joinpath(dump_file))
                        archive.extract(dump_file, target_dir)
        elif not self.use_extracted:
            raise CommandError(
                "Extracted dump already exists at {}. Delete it or use --use-extracted".format(target_dir))

        results = []
        filenames = []
        with futures.ThreadPoolExecutor(max_workers=len(export_meta_files)) as executor:
            for path in export_meta_files:
                results.append(executor.submit(_export_blobs, path, meta))

            for result in futures.as_completed(results):
{'fraction_non_alphanumeric': 0.08514518165240567, 'fraction_numerical': 0.00028054425585636136, 'mean_word_length': 3.1285466126230457, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '933740', 'n_tokens_mistral': 2015, 'n_tokens_neox': 1950, 'n_words': 553}