This is an automated email from the ASF dual-hosted git repository.

pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new ddc423bbd84 Fix mark-failed KeyError for removed-task TIs (#72338) 
(#72394)
ddc423bbd84 is described below

commit ddc423bbd848784d12b9a63c7b4ab564af9248ad
Author: Deepak Jain (DJ) <[email protected]>
AuthorDate: Mon Sep 7 03:08:20 2026 -0700

    Fix mark-failed KeyError for removed-task TIs (#72338) (#72394)
---
 airflow-core/src/airflow/api/common/mark_tasks.py     |  1 +
 airflow-core/tests/unit/api/common/test_mark_tasks.py | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py 
b/airflow-core/src/airflow/api/common/mark_tasks.py
index 1f597c0ee5f..31c81db0864 100644
--- a/airflow-core/src/airflow/api/common/mark_tasks.py
+++ b/airflow-core/src/airflow/api/common/mark_tasks.py
@@ -285,6 +285,7 @@ def _set_dag_run_terminal_state(
             select(TaskInstance).filter(
                 TaskInstance.dag_id == dag.dag_id,
                 TaskInstance.run_id == run_id,
+                TaskInstance.task_id.in_(task_ids),
                 or_(
                     TaskInstance.state.is_(None),
                     and_(
diff --git a/airflow-core/tests/unit/api/common/test_mark_tasks.py 
b/airflow-core/tests/unit/api/common/test_mark_tasks.py
index b16582ea58e..47278c74254 100644
--- a/airflow-core/tests/unit/api/common/test_mark_tasks.py
+++ b/airflow-core/tests/unit/api/common/test_mark_tasks.py
@@ -39,6 +39,25 @@ if TYPE_CHECKING:
 pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
 
 
+def test_set_dag_run_state_to_failed_ignores_removed_task_tis(dag_maker: 
DagMaker[SerializedDAG]):
+    with dag_maker("TEST_DAG_REMOVED_TASK"):
+        EmptyOperator(task_id="pending")
+        EmptyOperator(task_id="removed_task")
+    dr = dag_maker.create_dagrun()
+    dag_maker.session.flush()
+
+    with dag_maker("TEST_DAG_REMOVED_TASK") as dag_without_removed:
+        EmptyOperator(task_id="pending")
+
+    result: tuple[list[TaskInstance], list[TaskInstance]] = 
set_dag_run_state_to_failed(
+        dag=dag_without_removed, run_id=dr.run_id, commit=True, 
session=dag_maker.session
+    )
+    updated_tis, _ = result
+    assert len(updated_tis) == 1
+    assert updated_tis[0].task_id == "pending"
+    assert updated_tis[0].state == TaskInstanceState.SKIPPED
+
+
 def test_set_dag_run_state_to_failed(dag_maker: DagMaker[SerializedDAG]):
     with dag_maker("TEST_DAG_1") as dag:
         with EmptyOperator(task_id="teardown").as_teardown():

Reply via email to