pyinfra/pyinfra/payload_processing/payload.py

from dataclasses import dataclass
from itertools import chain
from operator import itemgetter
from typing import Union, Sized

from funcy import project

from pyinfra.config import Config
from pyinfra.utils.file_extension_parsing import make_file_extension_parser


@dataclass
class QueueMessagePayload:
    dossier_id: str
    file_id: str
    x_tenant_id: Union[str, None]

    target_file_extension: str
    response_file_extension: str

    target_file_type: Union[str, None]
    target_compression_type: Union[str, None]
    response_file_type: Union[str, None]
    response_compression_type: Union[str, None]

    target_file_name: str
    response_file_name: str

    processing_kwargs: dict


class QueueMessagePayloadParser:
    def __init__(self, file_extension_parser, allowed_processing_args=("operation",)):
        self.parse_file_extensions = file_extension_parser
        self.allowed_args = allowed_processing_args

    def __call__(self, payload: dict) -> QueueMessagePayload:
        """Translate the queue message payload to the internal QueueMessagePayload object."""
        return self._parse_queue_message_payload(payload)

    def _parse_queue_message_payload(self, payload: dict) -> QueueMessagePayload:
        dossier_id, file_id, target_file_extension, response_file_extension = itemgetter(
            "dossierId", "fileId", "targetFileExtension", "responseFileExtension"
        )(payload)
        x_tenant_id = payload.get("X-TENANT-ID")

        target_file_type, target_compression_type, response_file_type, response_compression_type = chain.from_iterable(
            map(self.parse_file_extensions, [target_file_extension, response_file_extension])
        )

        target_file_name = f"{dossier_id}/{file_id}.{target_file_extension}"
        response_file_name = f"{dossier_id}/{file_id}.{response_file_extension}"

        processing_kwargs = project(payload, self.allowed_args)

        return QueueMessagePayload(
            dossier_id=dossier_id,
            file_id=file_id,
            x_tenant_id=x_tenant_id,
            target_file_extension=target_file_extension,
            response_file_extension=response_file_extension,
            target_file_type=target_file_type,
            target_compression_type=target_compression_type,
            response_file_type=response_file_type,
            response_compression_type=response_compression_type,
            target_file_name=target_file_name,
            response_file_name=response_file_name,
            processing_kwargs=processing_kwargs,
        )


def get_queue_message_payload_parser(config: Config) -> QueueMessagePayloadParser:
    file_extension_parser = make_file_extension_parser(config.allowed_file_types, config.allowed_compression_types)
    return QueueMessagePayloadParser(file_extension_parser)


class QueueMessagePayloadFormatter:
    @staticmethod
    def format_service_processing_result_for_storage(
        queue_message_payload: QueueMessagePayload, service_processing_result: Sized
    ) -> dict:
        """Format the results of a processing function with the QueueMessagePayload for the storage upload."""
        return {
            "dossierId": queue_message_payload.dossier_id,
            "fileId": queue_message_payload.file_id,
            "targetFileExtension": queue_message_payload.target_file_extension,
            "responseFileExtension": queue_message_payload.response_file_extension,
            "data": service_processing_result,
        }

    @staticmethod
    def format_to_queue_message_response_body(queue_message_payload: QueueMessagePayload) -> dict:
        """Format QueueMessagePayload for the AMPQ response after processing."""
        return {"dossierId": queue_message_payload.dossier_id, "fileId": queue_message_payload.file_id}


def get_queue_message_payload_formatter() -> QueueMessagePayloadFormatter:
    return QueueMessagePayloadFormatter()