pyinfra/pyinfra/payload_processing/payload.py

from dataclasses import dataclass
from functools import singledispatch, partial
from funcy import project, complement
from itertools import chain
from operator import itemgetter
from typing import Union, Sized, Callable, List

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


@dataclass
class QueueMessagePayload:
    """Default one-to-one payload, where the message contains the absolute file paths for the target and response files,
    that have to be acquired from the storage."""

    target_file_path: str
    response_file_path: 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]

    x_tenant_id: Union[str, None]

    processing_kwargs: dict


@dataclass
class LegacyQueueMessagePayload(QueueMessagePayload):
    """Legacy one-to-one payload, where the message contains the dossier and file ids, and the file extensions that have
    to be used to construct the absolute file paths for the target and response files, that have to be acquired from the
    storage."""

    dossier_id: str
    file_id: str

    target_file_extension: str
    response_file_extension: str


class QueueMessagePayloadParser:
    def __init__(self, payload_matcher2parse_strategy: dict):
        self.payload_matcher2parse_strategy = payload_matcher2parse_strategy

    def __call__(self, payload: dict) -> QueueMessagePayload:
        for payload_matcher, parse_strategy in self.payload_matcher2parse_strategy.items():
            if payload_matcher(payload):
                return parse_strategy(payload)


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

    payload_matcher2parse_strategy = get_payload_matcher2parse_strategy(
        file_extension_parser, config.allowed_processing_parameters
    )

    return QueueMessagePayloadParser(payload_matcher2parse_strategy)


def get_payload_matcher2parse_strategy(parse_file_extensions: Callable, allowed_processing_parameters: List[str]):
    return {
        is_legacy_payload: partial(
            parse_legacy_queue_message_payload,
            parse_file_extensions=parse_file_extensions,
            allowed_processing_parameters=allowed_processing_parameters,
        ),
        complement(is_legacy_payload): partial(
            parse_queue_message_payload,
            parse_file_extensions=parse_file_extensions,
            allowed_processing_parameters=allowed_processing_parameters,
        ),
    }


def is_legacy_payload(payload: dict) -> bool:
    return {"dossierId", "fileId", "targetFileExtension", "responseFileExtension"}.issubset(payload.keys())


def parse_queue_message_payload(
    payload: dict,
    parse_file_extensions: Callable,
    allowed_processing_parameters: List[str],
) -> QueueMessagePayload:
    target_file_path, response_file_path = itemgetter("targetFilePath", "responseFilePath")(payload)

    target_file_type, target_compression_type, response_file_type, response_compression_type = chain.from_iterable(
        map(parse_file_extensions, [target_file_path, response_file_path])
    )

    x_tenant_id = payload.get("X-TENANT-ID")

    processing_kwargs = project(payload, allowed_processing_parameters)

    return QueueMessagePayload(
        target_file_path=target_file_path,
        response_file_path=response_file_path,
        target_file_type=target_file_type,
        target_compression_type=target_compression_type,
        response_file_type=response_file_type,
        response_compression_type=response_compression_type,
        x_tenant_id=x_tenant_id,
        processing_kwargs=processing_kwargs,
    )


def parse_legacy_queue_message_payload(
    payload: dict,
    parse_file_extensions: Callable,
    allowed_processing_parameters: List[str],
) -> LegacyQueueMessagePayload:
    dossier_id, file_id, target_file_extension, response_file_extension = itemgetter(
        "dossierId", "fileId", "targetFileExtension", "responseFileExtension"
    )(payload)

    target_file_path = f"{dossier_id}/{file_id}.{target_file_extension}"
    response_file_path = f"{dossier_id}/{file_id}.{response_file_extension}"

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

    x_tenant_id = payload.get("X-TENANT-ID")

    processing_kwargs = project(payload, allowed_processing_parameters)

    return LegacyQueueMessagePayload(
        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_path=target_file_path,
        response_file_path=response_file_path,
        processing_kwargs=processing_kwargs,
    )


@singledispatch
def format_service_processing_result_for_storage(payload: QueueMessagePayload, result: Sized) -> dict:
    raise NotImplementedError("Unsupported payload type")


@format_service_processing_result_for_storage.register(LegacyQueueMessagePayload)
def _(payload: LegacyQueueMessagePayload, result: Sized) -> dict:
    processing_kwargs = payload.processing_kwargs or {}
    x_tenant_id = {"X-TENANT-ID": payload.x_tenant_id} if payload.x_tenant_id else {}
    return {
        "dossierId": payload.dossier_id,
        "fileId": payload.file_id,
        "targetFileExtension": payload.target_file_extension,
        "responseFileExtension": payload.response_file_extension,
        **x_tenant_id,
        **processing_kwargs,
        "data": result,
    }


@format_service_processing_result_for_storage.register(QueueMessagePayload)
def _(payload: QueueMessagePayload, result: Sized) -> dict:
    processing_kwargs = payload.processing_kwargs or {}
    x_tenant_id = {"X-TENANT-ID": payload.x_tenant_id} if payload.x_tenant_id else {}
    return {
        "targetFilePath": payload.target_file_path,
        "responseFilePath": payload.response_file_path,
        **x_tenant_id,
        **processing_kwargs,
        "data": result,
    }


@singledispatch
def format_to_queue_message_response_body(queue_message_payload: QueueMessagePayload) -> dict:
    raise NotImplementedError("Unsupported payload type")


@format_to_queue_message_response_body.register(LegacyQueueMessagePayload)
def _(payload: LegacyQueueMessagePayload) -> dict:
    processing_kwargs = payload.processing_kwargs or {}
    x_tenant_id = {"X-TENANT-ID": payload.x_tenant_id} if payload.x_tenant_id else {}
    return {"dossierId": payload.dossier_id, "fileId": payload.file_id, **x_tenant_id, **processing_kwargs}


@format_to_queue_message_response_body.register(QueueMessagePayload)
def _(payload: QueueMessagePayload) -> dict:
    processing_kwargs = payload.processing_kwargs or {}
    x_tenant_id = {"X-TENANT-ID": payload.x_tenant_id} if payload.x_tenant_id else {}
    return {
        "targetFilePath": payload.target_file_path,
        "responseFilePath": payload.response_file_path,
        **x_tenant_id,
        **processing_kwargs,
    }