<commit_before>"""
This package contains classes and function for processing ExportInstances using
multiple processes.

To rebuild an export run the following:

    rebuild_export_mutiprocess(export_instance_id, num_processes, page_size)

You can also use the MultiprocessExporter class to have more control over the process.
See the 'process_skipped_pages' management command for an example.

The export works as follows:
  * Dump raw docs from ES into files of size N docs
  * Once each file is complete add it to a multiprocessing Queue
  * Pool of X processes listen to queue and process the dump file
  * Results returned back to the main process
    * Unsuccessful results can be retried
  * Add successful pages to final ZIP archive
  * Add raw data dumps for unsuccessful pages to final ZIP archive
"""
import gzip
import json
import logging
import multiprocessing
import os
import tempfile
import time
import zipfile
from Queue import Empty
from collections import namedtuple
from datetime import timedelta

from corehq.apps.export.dbaccessors import get_properly_wrapped_export_instance
from corehq.apps.export.export import (
    ExportFile, get_export_writer, save_export_payload, get_export_size, get_export_documents)
from corehq.apps.export.export import write_export_instance
from corehq.elastic import ScanResult
from corehq.util.files import safe_filename
from couchexport.export import get_writer
from couchexport.writers import ZippedExportWriter

TEMP_FILE_PREFIX = 'cchq_export_dump_'

UNPROCESSED_PAGES_DIR = 'unprocessed'

logger = logging.getLogger(__name__)


ProgressValue = namedtuple('ProgressValue', 'page progress total')


class BaseResult(object):
    success = False

    def __init__(self, page_number, page_path, page_size):
        self.page = page_number
        self.path = page_path
        self.page_size = page_size


class SuccessResult(BaseResult):
    success = True


class RetryResult(BaseResult):
    def __init__(self, page_number, page_path, page_size, retry_count):
        super(RetryResult, self).__init__(page_number, page_path, page_size)
        self.retry_count = retry_count


class QueuedResult(RetryResult):
    def __init__(self, async_result, page_number, page_path, page_size, retry_count):
        super(QueuedResult, self).__init__(page_number, page_path, page_size, retry_count)
        self.async_result = async_result


class OutputPaginator(object):
    """Helper class to paginate raw export output"""
    def __init__(self, export_id):
        self.export_id = export_id
        self.page = 0
        self.page_size = 0
        self.file = None

    def __enter__(self):
        self._new_file()

    def _new_file(self):
        if self.file:
            self.file.close()
        prefix = '{}{}_'.format(TEMP_FILE_PREFIX, self.export_id)
        fileobj = tempfile.NamedTemporaryFile(prefix=prefix, mode='wb', delete=False)
        self.path = fileobj.name
        self.file = gzip.GzipFile(fileobj=fileobj)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()
        if exc_type is not None:
            os.remove(self.path)
        self.file = None
        self.path = None

    def next_page(self):
        self.page += 1
        self.page_size = 0
        self._new_file()

    def write(self, doc):
        self.page_size += 1
        self.file.write('{}\n'.format(json.dumps(doc)))

    def get_result(self):
        return RetryResult(self.page, self.path, self.page_size, 0)


def rebuild_export_mutiprocess(export_id, num_processes, page_size=100000):
    assert num_processes > 0

    def _log_page_dumped(paginator):
        logger.info('  Dump page {} complete: {} docs'.format(paginator.page, paginator.page_size))

    export_instance = get_properly_wrapped_export_instance(export_id)
    filters = export_instance.get_filters()
    total_docs = get_export_size(export_instance, filters)

    logger.info('Starting data dump of {} docs'.format(total_docs))
    exporter = MultiprocessExporter(export_instance, total_docs, num_processes)
    paginator = OutputPaginator(export_id)
    with exporter, paginator:
        for index, doc in enumerate(get_export_documents(export_instance, filters)):
            paginator.write(doc)
            if paginator.page_size == page_size:
                _log_page_dumped(paginator)
                exporter.process_page(paginator.get_result())
                paginator.next_page()
        if paginator.page_size:
            _log_page_dumped(paginator)
            exporter.process_page(paginator.get_result())

    exporter.wait_till_completion()


def run_export_with_logging(export_instance, page_number, dump_path, doc_count, attempts):
    """Log any exceptions here since logging on the other side of the process queue
    won't show the traceback
    """<commit_msg>remove unused enumerated<commit_after>"""
This package contains classes and function for processing ExportInstances using
multiple processes.

To rebuild an export run the following:

    rebuild_export_mutiprocess(export_instance_id, num_processes, page_size)

You can also use the MultiprocessExporter class to have more control over the process.
See the 'process_skipped_pages' management command for an example.

The export works as follows:
  * Dump raw docs from ES into files of size N docs
  * Once each file is complete add it to a multiprocessing Queue
  * Pool of X processes listen to queue and process the dump file
  * Results returned back to the main process
    * Unsuccessful results can be retried
  * Add successful pages to final ZIP archive
  * Add raw data dumps for unsuccessful pages to final ZIP archive
"""
import gzip
import json
import logging
import multiprocessing
import os
import tempfile
import time
import zipfile
from Queue import Empty
from collections import namedtuple
from datetime import timedelta

from corehq.apps.export.dbaccessors import get_properly_wrapped_export_instance
from corehq.apps.export.export import (
    ExportFile, get_export_writer, save_export_payload, get_export_size, get_export_documents)
from corehq.apps.export.export import write_export_instance
from corehq.elastic import ScanResult
from corehq.util.files import safe_filename
from couchexport.export import get_writer
from couchexport.writers import ZippedExportWriter

TEMP_FILE_PREFIX = 'cchq_export_dump_'

UNPROCESSED_PAGES_DIR = 'unprocessed'

logger = logging.getLogger(__name__)


ProgressValue = namedtuple('ProgressValue', 'page progress total')


class BaseResult(object):
    success = False

    def __init__(self, page_number, page_path, page_size):
        self.page = page_number
        self.path = page_path
        self.page_size = page_size


class SuccessResult(BaseResult):
    success = True


class RetryResult(BaseResult):
    def __init__(self, page_number, page_path, page_size, retry_count):
        super(RetryResult, self).__init__(page_number, page_path, page_size)
        self.retry_count = retry_count


class QueuedResult(RetryResult):
    def __init__(self, async_result, page_number, page_path, page_size, retry_count):
        super(QueuedResult, self).__init__(page_number, page_path, page_size, retry_count)
        self.async_result = async_result


class OutputPaginator(object):
    """Helper class to paginate raw export output"""
    def __init__(self, export_id):
        self.export_id = export_id
        self.page = 0
        self.page_size = 0
        self.file = None

    def __enter__(self):
        self._new_file()

    def _new_file(self):
        if self.file:
            self.file.close()
        prefix = '{}{}_'.format(TEMP_FILE_PREFIX, self.export_id)
        fileobj = tempfile.NamedTemporaryFile(prefix=prefix, mode='wb', delete=False)
        self.path = fileobj.name
        self.file = gzip.GzipFile(fileobj=fileobj)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()
        if exc_type is not None:
            os.remove(self.path)
        self.file = None
        self.path = None

    def next_page(self):
        self.page += 1
        self.page_size = 0
        self._new_file()

    def write(self, doc):
        self.page_size += 1
        self.file.write('{}\n'.format(json.dumps(doc)))

    def get_result(self):
        return RetryResult(self.page, self.path, self.page_size, 0)


def rebuild_export_mutiprocess(export_id, num_processes, page_size=100000):
    assert num_processes > 0

    def _log_page_dumped(paginator):
        logger.info('  Dump page {} complete: {} docs'.format(paginator.page, paginator.page_size))

    export_instance = get_properly_wrapped_export_instance(export_id)
    filters = export_instance.get_filters()
    total_docs = get_export_size(export_instance, filters)

    logger.info('Starting data dump of {} docs'.format(total_docs))
    exporter = MultiprocessExporter(export_instance, total_docs, num_processes)
    paginator = OutputPaginator(export_id)
    with exporter, paginator:
        for doc in get_export_documents(export_instance, filters):
            paginator.write(doc)
            if paginator.page_size == page_size:
                _log_page_dumped(paginator)
                exporter.process_page(paginator.get_result())
                paginator.next_page()
        if paginator.page_size:
            _log_page_dumped(paginator)
            exporter.process_page(paginator.get_result())

    exporter.wait_till_completion()


def run_export_with_logging(export_instance, page_number, dump_path, doc_count, attempts):
    """Log any exceptions here since logging on the other side of the process queue
    won't show the traceback
    """
{'fraction_non_alphanumeric': 0.07579056505961639, 'fraction_numerical': 0.002695697252462416, 'mean_word_length': 3.687074829931973, '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': '3498751', 'n_tokens_mistral': 2937, 'n_tokens_neox': 2824, 'n_words': 775}