diff --git a/tests/utilities/aws_lambda/test_task_handler.py b/tests/utilities/aws_lambda/test_task_handler.py index 974a47966..b986c1c5e 100644 --- a/tests/utilities/aws_lambda/test_task_handler.py +++ b/tests/utilities/aws_lambda/test_task_handler.py @@ -5,7 +5,7 @@ from typing import Any from uuid import UUID import pytest -from sqlalchemy import Engine +from sqlalchemy import Engine, text from sqlmodel import Session from domain.tasks.tasks import Source @@ -148,14 +148,9 @@ def test_task_handler_attaches_to_a_supplied_task_and_subtask( {"task_id": str(task.id), "subtask_id": str(subtask.id)} ] assert harness.subtasks.get(subtask.id).status.value == "complete" - with Session(db_engine) as session: - from sqlmodel import select - - from backend.app.db.models.tasks import SubTask as SubTaskRow - from backend.app.db.models.tasks import Task as TaskRow - - assert len(session.exec(select(TaskRow)).all()) == 1 - assert len(session.exec(select(SubTaskRow)).all()) == 1 + with db_engine.connect() as conn: + assert conn.execute(text("SELECT count(*) FROM tasks")).scalar_one() == 1 + assert conn.execute(text("SELECT count(*) FROM sub_task")).scalar_one() == 1 def test_task_handler_leaves_cloudwatch_url_unset_outside_lambda( diff --git a/utilities/aws_lambda/task_handler.py b/utilities/aws_lambda/task_handler.py index f2064f573..695153add 100644 --- a/utilities/aws_lambda/task_handler.py +++ b/utilities/aws_lambda/task_handler.py @@ -9,6 +9,7 @@ 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 @@ -50,39 +51,47 @@ def task_handler( 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, - ) + # 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) + 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)} + {"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), + subtask_id, + work=lambda: func(body, context, orchestrator, task_id), cloud_logs_url=cloud_logs_url, ) else: orchestrator.run_subtask( - subtask.id, + subtask_id, work=lambda: func(body, context), cloud_logs_url=cloud_logs_url, ) except Exception: logger.exception( - "subtask failed (task_source=%s source_id=%s)", + "subtask failed (task_source=%s subtask_id=%s)", task_source, - source_id, + subtask_id, ) if "Records" in event: message_id = record.get("messageId", "") @@ -99,6 +108,16 @@ def task_handler( 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):