pyinfra/pyinfra/payload_processing/processor.py

import logging
from dataclasses import asdict
from functools import partial
from typing import Callable, Union, List

from funcy import compose

from pyinfra.config import get_config, Config
from pyinfra.payload_processing.payload import (
    read_queue_message_payload,
    format_service_processing_result_for_storage,
    format_to_queue_message_response_body,
)
from pyinfra.storage import get_storage
from pyinfra.storage.storage import make_downloader, make_uploader
from pyinfra.storage.storages.interface import Storage

logger = logging.getLogger()
logger.setLevel(get_config().logging_level_root)


class PayloadProcessor:
    def __init__(self, storage: Storage, bucket: str, data_processor: Callable):
        """Wraps an analysis function specified by a service (e.g. NER service) in pre- and post-processing steps.

        Args:
            storage: The storage to use for downloading and uploading files
            bucket: The bucket to use for downloading and uploading files
            data_processor: The analysis function to be called with the downloaded file
        """
        self.process_data = data_processor

        self.partial_download_fn = partial(make_downloader, storage, bucket)
        self.partial_upload_fn = partial(make_uploader, storage, bucket)

    def __call__(self, queue_message_payload: dict) -> dict:
        """Processes a queue message payload.

        The steps executed are:
            1. Download the file specified in the message payload from the storage
            2. Process the file with the analysis function
            3. Upload the result to the storage
            4. Return the payload for a response queue message

        Args:
            queue_message_payload: The payload of a queue message. The payload is expected to be a dict with the
                following keys: dossierId, fileId, targetFileExtension, responseFileExtension

        Returns:
            The payload for a response queue message. The payload is a dict with the following keys: dossierId, fileId
        """
        return self._process(queue_message_payload)

    def _process(self, queue_message_payload: dict) -> dict:
        payload = read_queue_message_payload(queue_message_payload)
        logger.info(f"Processing {asdict(payload)} ...")

        download_file_to_process = self.partial_download_fn(payload.target_file_type, payload.target_compression_type)
        upload_processing_result = self.partial_upload_fn(payload.response_file_type, payload.response_compression_type)
        format_result_for_storage = partial(format_service_processing_result_for_storage, payload)

        processing_pipeline = compose(format_result_for_storage, self.process_data, download_file_to_process)

        result: List[dict] = processing_pipeline(payload.target_file_name)

        upload_processing_result(payload.response_file_name, result)

        return format_to_queue_message_response_body(payload)


def make_payload_processor(data_processor: Callable, config: Union[None, Config]) -> PayloadProcessor:
    """Produces payload processor for queue manager."""
    config = config or get_config()

    storage: Storage = get_storage(config)
    bucket: str = config.storage_bucket

    data_processor = compose(list, data_processor)

    return PayloadProcessor(storage, bucket, data_processor)