(airflow) branch main updated: Fix Callback exception when a removed task is the last one in the task instance list (#36693)

[ep...@apache.org]

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

ephraimanierobi 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 8c1c09bab3 Fix Callback exception when a removed task is the last one in the task instance list (#36693)
8c1c09bab3 is described below

commit 8c1c09bab34be8234d9e152ed7e4c9b925c08459
Author: Joao Amaral <72...@users.noreply.github.com>
AuthorDate: Mon Jan 15 07:50:10 2024 -0300

    Fix Callback exception when a removed task is the last one in the task instance list (#36693)
    
    * Fix Callback exception when a removed task is the last one in the task instance list
    
    * Add test_dag_handle_callback_with_removed_task
    
    * Remove extra break line
    
    * Merge TIs filters
    
    * Fix static check
    
    * Revert changes
---
 airflow/models/dag.py    |  2 ++
 tests/models/test_dag.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index 859d78b0ab..4e70b87817 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -1466,6 +1466,8 @@ class DAG(LoggingMixin):
             # context for the callback.
             if dag.partial:
                 tis = [ti for ti in tis if not ti.state == State.NONE]
+            # filter out removed tasks
+            tis = [ti for ti in tis if ti.state != TaskInstanceState.REMOVED]
             ti = tis[-1]  # get first TaskInstance of DagRun
             ti.task = dag.get_task(ti.task_id)
             context = ti.get_template_context(session=session)
diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py
index de7fa87d5a..f367b00abe 100644
--- a/tests/models/test_dag.py
+++ b/tests/models/test_dag.py
@@ -1512,6 +1512,34 @@ class TestDag:
         dag.clear()
         self._clean_up(dag_id)
 
+    def test_dag_handle_callback_with_removed_task(self, dag_maker, session):
+        """
+        Tests avoid crashes when a removed task is the last one in the list of task instance
+        """
+        dag_id = "test_dag_callback_with_removed_task"
+        mock_callback = mock.MagicMock()
+        with DAG(
+            dag_id=dag_id,
+            on_success_callback=mock_callback,
+            on_failure_callback=mock_callback,
+        ) as dag:
+            EmptyOperator(task_id="faketastic")
+            task_removed = EmptyOperator(task_id="removed_task")
+
+        with create_session() as session:
+            dag_run = dag.create_dagrun(State.RUNNING, TEST_DATE, run_type=DagRunType.MANUAL, session=session)
+            dag._remove_task(task_removed.task_id)
+            tis = dag_run.get_task_instances(session=session)
+            tis[-1].state = TaskInstanceState.REMOVED
+            assert dag_run.get_task_instance(task_removed.task_id).state == TaskInstanceState.REMOVED
+
+            # should not raise any exception
+            dag.handle_callback(dag_run, success=True)
+            dag.handle_callback(dag_run, success=False)
+
+        dag.clear()
+        self._clean_up(dag_id)
+
     def test_next_dagrun_after_fake_scheduled_previous(self):
         """
         Test scheduling a dag where there is a prior DagRun