Model/utilities/aws_lambda/task_handler.py

"""@task_handler decorator for Lambdas that own the entire pipeline.

Translates an AWS Lambda invocation (SQS-shaped or direct) into
TaskOrchestrator.create_task_with_subtask(...) + run_subtask(...).
"""

import json
import logging
from contextlib import AbstractContextManager
from functools import wraps
from typing import Any, Callable, Optional, cast
from uuid import UUID

from utilities.aws_lambda.cloud_logs import cloudwatch_url
from utilities.aws_lambda.default_orchestrator import default_orchestrator
from domain.tasks.subtasks import SubTaskFailure
from domain.tasks.tasks import Source
from orchestration.task_orchestrator import TaskOrchestrator

logger = logging.getLogger(__name__)

OrchestratorCM = Callable[[], AbstractContextManager[TaskOrchestrator]]


class NonRetriableTaskError(Exception):
    """A failure that must not be retried by SQS redelivery.

    The SubTask is marked failed (the failure record is the visibility
    surface) but the message is not returned to the queue, so the work is
    never re-run. Raise it for failures where a retry could duplicate an
    irreversible side effect — e.g. a job already logged in OpenHousing
    whose HubSpot write-back failed.
    """


def task_handler(
    *,
    task_source: str,
    source: Source,
    orchestrator_cm: Optional[OrchestratorCM] = None,
    pass_task_orchestrator: bool = False,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Run the wrapped function as the body of a freshly-created Task + SubTask.

    For each record, creates a new Task + initial SubTask, then runs the
    wrapped function inside orchestrator.run_subtask(...). `source_id` is
    read from body[source.value] (silent None if absent — preserved from
    legacy ADR-0001).

    Records-style events use SQS partial-batch-failure semantics: individual
    failures are reported via {"batchItemFailures": [...]} rather than
    propagating. Direct invocations re-raise.
    """
    factory = orchestrator_cm or default_orchestrator

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        def wrapper(event: dict[str, Any], context: Any) -> Any:
            cloud_logs_url = cloudwatch_url()
            with factory() as orchestrator:
                task_ids: list[dict[str, str]] = []
                failures: list[dict[str, Any]] = []

                for record in _records(event):
                    body = _parse_body(record)

                    # Attach mode (ADR-0055): a Modelling Run batch carries the
                    # app-owned task_id and its distributor-pre-created
                    # subtask_id — run under those; create nothing.
                    supplied = _supplied_ids(body)
                    source_id: Optional[str] = None
                    if supplied is not None:
                        task_id, subtask_id = supplied
                    else:
                        raw_source_id = body.get(source.value)
                        source_id = (
                            str(raw_source_id) if raw_source_id is not None else None
                        )
                        task, subtask = orchestrator.create_task_with_subtask(
                            task_source=task_source,
                            inputs=body,
                            source=source,
                            source_id=source_id,
                        )
                        task_id, subtask_id = task.id, subtask.id
                    task_ids.append(
                        {"task_id": str(task_id), "subtask_id": str(subtask_id)}
                    )

                    try:
                        if pass_task_orchestrator:
                            orchestrator.run_subtask(
                                subtask_id,
                                work=lambda: func(body, context, orchestrator, task_id),
                                cloud_logs_url=cloud_logs_url,
                            )
                        else:
                            orchestrator.run_subtask(
                                subtask_id,
                                work=lambda: func(body, context),
                                cloud_logs_url=cloud_logs_url,
                            )
                    except NonRetriableTaskError:
                        # Failed, but never redelivered: the SubTask's failure
                        # record is the only follow-up surface.
                        logger.exception(
                            "subtask failed non-retriably "
                            "(task_source=%s source_id=%s)",
                            task_source,
                            source_id,
                        )
                        if "Records" not in event:
                            raise
                    except SubTaskFailure as recorded:
                        # A recorded failure (ADR-0055): run_subtask has already
                        # failed the SubTask with the structured details — that
                        # record IS the outcome. Never returned to SQS for retry;
                        # recovery is a deliberate re-send of the sub_task's
                        # own inputs.
                        logger.warning(
                            "subtask recorded failure, not retrying "
                            "(task_source=%s subtask_id=%s): %s",
                            task_source,
                            subtask_id,
                            recorded,
                        )
                    except Exception:
                        logger.exception(
                            "subtask failed (task_source=%s subtask_id=%s)",
                            task_source,
                            subtask_id,
                        )
                        if "Records" in event:
                            message_id = record.get("messageId", "")
                            failures.append({"itemIdentifier": message_id})
                        else:
                            raise

                if "Records" in event:
                    return {"batchItemFailures": failures, "tasks": task_ids}
                return task_ids

        return wrapper

    return decorator


def _supplied_ids(body: dict[str, Any]) -> Optional[tuple[UUID, UUID]]:
    """The (task_id, subtask_id) an attach-mode message carries, or None for
    the classic create-your-own-task message."""
    raw_task_id = body.get("task_id")
    raw_subtask_id = body.get("subtask_id")
    if raw_task_id is None or raw_subtask_id is None:
        return None
    return UUID(str(raw_task_id)), UUID(str(raw_subtask_id))


def _parse_body(record: dict[str, Any]) -> dict[str, Any]:
    raw = record.get("body", record)
    if isinstance(raw, str):
        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError:
            return {}
        return cast(dict[str, Any], parsed) if isinstance(parsed, dict) else {}
    if isinstance(raw, dict):
        return cast(dict[str, Any], raw)
    return {}


def _records(event: dict[str, Any]) -> list[dict[str, Any]]:
    raw_records = event.get("Records")
    if isinstance(raw_records, list):
        return [r for r in cast(list[Any], raw_records) if isinstance(r, dict)]
    return [event]