From 550ff65678bf88c0378b515e0a1adb4ad7935cf1 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 14 Jan 2022 14:10:17 +0100 Subject: [PATCH] Fix failing main (#20871) When I merged #18724 the jobs ran successfully but it's now failing in main. This PR fixes it Co-authored-by: Kaxil Naik Co-authored-by: Ash Berlin-Taylor GitOrigin-RevId: b1fffe7b2673f5aaad72304d999f4a55b2cc2edc --- airflow/api_connexion/endpoints/task_instance_endpoint.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/api_connexion/endpoints/task_instance_endpoint.py b/airflow/api_connexion/endpoints/task_instance_endpoint.py index 1cef6e9ab33..cbcc189e9c2 100644 --- a/airflow/api_connexion/endpoints/task_instance_endpoint.py +++ b/airflow/api_connexion/endpoints/task_instance_endpoint.py @@ -326,7 +326,9 @@ def post_set_task_instances_state(*, dag_id: str, session: Session = NEW_SESSION detail=f"Task instance not found for task {task_id!r} on execution_date {execution_date}" ) - if run_id and not session.query(TI).get({'task_id': task_id, 'dag_id': dag_id, 'run_id': run_id}): + if run_id and not session.query(TI).get( + {'task_id': task_id, 'dag_id': dag_id, 'run_id': run_id, 'map_index': -1} + ): error_message = f"Task instance not found for task {task_id!r} on DAG run with ID {run_id!r}" raise NotFound(detail=error_message)