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():