diff --git a/applications/audit_generator/handler.py b/applications/audit_generator/handler.py index 76f53538..72358143 100644 --- a/applications/audit_generator/handler.py +++ b/applications/audit_generator/handler.py @@ -16,7 +16,7 @@ from orchestration.audit_generator_unit_of_work import AuditGeneratorUnitOfWork from utilities.aws_lambda.subtask_handler import subtask_handler -@subtask_handler() +@subtask_handler(task_orchestrator=False) def handler(body: dict[str, Any], context: Any) -> None: trigger = AuditGeneratorTriggerRequest.model_validate(body) diff --git a/utilities/aws_lambda/subtask_handler.py b/utilities/aws_lambda/subtask_handler.py index 42c9e07a..c53b59ca 100644 --- a/utilities/aws_lambda/subtask_handler.py +++ b/utilities/aws_lambda/subtask_handler.py @@ -26,6 +26,7 @@ OrchestratorCM = Callable[[], AbstractContextManager[TaskOrchestrator]] def subtask_handler( *, orchestrator_cm: Optional[OrchestratorCM] = None, + task_orchestrator: Optional[bool] = True, ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: """Run the wrapped function as the body of an existing SubTask. @@ -38,7 +39,7 @@ def subtask_handler( factory = orchestrator_cm or default_orchestrator def decorator(func: Callable[..., Any]) -> Callable[..., Any]: - _wants_orchestrator = len(inspect.signature(func).parameters) >= 3 + _wants_orchestrator = task_orchestrator @wraps(func) def wrapper(event: dict[str, Any], context: Any) -> None: @@ -48,6 +49,7 @@ def subtask_handler( body = _parse_body(record) trigger = SubtaskTriggerBody.model_validate(body) logger.info("Running subtask %s", trigger.sub_task_id) + def _work_with( _body: dict[str, Any] = body, _o: TaskOrchestrator = orchestrator, @@ -67,9 +69,7 @@ def subtask_handler( cloud_logs_url=cloud_logs_url, ) except Exception: - logger.exception( - "Subtask %s failed", trigger.sub_task_id - ) + logger.exception("Subtask %s failed", trigger.sub_task_id) raise logger.info("Subtask %s completed", trigger.sub_task_id)