wyndham-ARR/arr_processing/registration.py

"""ARR-side registration of authenticated remote file handles into private storage."""

from __future__ import annotations

import hashlib
import json
import os
import stat
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Mapping, Protocol

from arr_ingestion.contracts import (
    ARTIFACT_SIZE_LIMITS,
    DELIVERY_SCHEMA_VERSION,
    ArtifactRef,
    DeliveryEnvelope,
)
from arr_processing.contracts import ProcessingResult, ResultArtifactHandle
from arr_processing.errors import ProcessingError, ProcessingTransportError
from arr_storage.contracts import StoredObject
from arr_storage.store import ManagedObjectStore


class RemoteFilePort(Protocol):
    """Official runtime/file API seam; it receives only one opaque file handle."""

    def materialize(self, file_handle: str, destination: Path, max_bytes: int) -> None:
        ...


@dataclass(frozen=True)
class RegisteredDelivery:
    envelope: DeliveryEnvelope
    raw_envelope: bytes
    output_objects: Mapping[str, StoredObject]


class ProcessingOutputRegistrar:
    def __init__(
        self,
        remote_files: RemoteFilePort,
        object_store: ManagedObjectStore,
    ) -> None:
        self._remote_files = remote_files
        self._object_store = object_store

    def register(
        self,
        result: ProcessingResult,
        source: ArtifactRef,
    ) -> RegisteredDelivery:
        if source.role != "source_xml":
            raise ProcessingError(
                "PROCESSING_SOURCE_INVALID", "processing source artifact is invalid"
            )
        verified_source = self._object_store.inspect_committed(
            source.object_key,
            source.original_filename,
        ).to_artifact_ref()
        if verified_source != source:
            raise ProcessingError(
                "PROCESSING_SOURCE_INVALID", "processing source identity does not match"
            )
        handles = [
            artifact.file_handle
            for artifact in result.artifacts.values()
            if artifact is not None
        ]
        if len(handles) != len(set(handles)):
            raise ProcessingError(
                "PROCESSING_ARTIFACT_INVALID", "processing artifact handles are not unique"
            )

        stored: Dict[str, StoredObject] = {}
        with tempfile.TemporaryDirectory(prefix="arr-processing-output-") as temporary:
            root = Path(temporary)
            for role in sorted(result.artifacts):
                artifact = result.artifacts[role]
                if artifact is None:
                    continue
                destination = root / role / artifact.original_filename
                destination.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
                try:
                    self._remote_files.materialize(
                        artifact.file_handle,
                        destination,
                        ARTIFACT_SIZE_LIMITS[role],
                    )
                except ProcessingTransportError as error:
                    raise ProcessingError(
                        error.code,
                        "remote processing artifact could not be fetched",
                        retryable=error.retryable,
                    ) from None
                self._verify_materialized(destination, artifact)
                stored[role] = self._object_store.upload_committed(
                    job_id=result.job_id,
                    attempt_no=result.attempt_no,
                    role=role,
                    source=destination,
                    original_filename=artifact.original_filename,
                    expected_sha256=artifact.sha256,
                    expected_byte_size=artifact.byte_size,
                )

        artifacts: Dict[str, object] = {
            "source_xml": source.to_dict(),
            "daily_report": None,
            "result_json": None,
            "structured_result_json": None,
            "exception_report": None,
        }
        for role, value in stored.items():
            artifacts[role] = value.to_artifact_ref().to_dict()
        payload = {
            "delivery_schema_version": DELIVERY_SCHEMA_VERSION,
            "delivery_id": result.delivery_id,
            "job_id": result.job_id,
            "attempt_no": result.attempt_no,
            "status": result.status,
            "processor_version": result.processor_version,
            "rule_set_sha256": result.rule_set_sha256,
            "result_schema_version": result.result_schema_version,
            "business_date": (
                result.business_date.isoformat() if result.business_date is not None else None
            ),
            "artifacts": artifacts,
        }
        envelope = DeliveryEnvelope.from_dict(payload)
        raw = (
            json.dumps(
                envelope.to_dict(),
                ensure_ascii=False,
                allow_nan=False,
                sort_keys=True,
                separators=(",", ":"),
            )
            + "\n"
        ).encode("utf-8")
        return RegisteredDelivery(envelope, raw, stored)

    @staticmethod
    def _verify_materialized(destination: Path, artifact: ResultArtifactHandle) -> None:
        try:
            file_stat = destination.lstat()
        except OSError:
            raise ProcessingError(
                "PROCESSING_ARTIFACT_NOT_FOUND", "remote processing artifact is unavailable"
            ) from None
        if destination.is_symlink() or not stat.S_ISREG(file_stat.st_mode):
            raise ProcessingError(
                "PROCESSING_ARTIFACT_INVALID", "remote processing artifact is not a regular file"
            )
        if file_stat.st_size != artifact.byte_size:
            raise ProcessingError(
                "PROCESSING_ARTIFACT_HASH_MISMATCH", "remote artifact identity does not match"
            )
        digest = hashlib.sha256()
        descriptor = -1
        try:
            descriptor = os.open(
                destination,
                os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
            )
            with os.fdopen(descriptor, "rb") as source:
                descriptor = -1
                while True:
                    chunk = source.read(1024 * 1024)
                    if not chunk:
                        break
                    digest.update(chunk)
        except OSError:
            raise ProcessingError(
                "PROCESSING_ARTIFACT_INVALID", "remote processing artifact could not be read"
            ) from None
        finally:
            if descriptor >= 0:
                os.close(descriptor)
        if digest.hexdigest() != artifact.sha256:
            raise ProcessingError(
                "PROCESSING_ARTIFACT_HASH_MISMATCH", "remote artifact identity does not match"
            )