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 utilities.aws_lambda.cloud_logs import cloudwatch_url
from utilities.aws_lambda.default_orchestrator import default_orchestrator
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)
                    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_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 Exception:
                        logger.exception(
                            "subtask failed (task_source=%s source_id=%s)",
                            task_source,
                            source_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 _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]