From 1341174b8869bc5cef892d6d25bcbf2ef9aab118 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Tue, 14 Jul 2026 17:08:29 +0200 Subject: [PATCH 1/2] Call listeners for running task instance when a Dag Run state is manually set --- .../src/airflow/api/common/mark_tasks.py | 35 ++++++---- .../core_api/services/public/dag_run.py | 17 ++++- .../tests/unit/api/common/test_mark_tasks.py | 9 ++- .../core_api/routes/public/test_dag_run.py | 66 ++++++++++++++++++- 4 files changed, 110 insertions(+), 17 deletions(-) diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py b/airflow-core/src/airflow/api/common/mark_tasks.py index 5274d8c706e04..1e81a6b174f1d 100644 --- a/airflow-core/src/airflow/api/common/mark_tasks.py +++ b/airflow-core/src/airflow/api/common/mark_tasks.py @@ -223,7 +223,7 @@ def _set_dag_run_terminal_state( ti_state: TaskInstanceState, commit: bool, session: SASession, -) -> list[TaskInstance]: +) -> tuple[list[TaskInstance], list[TaskInstance]]: """ Set the dag run's state to the given terminal state. @@ -237,11 +237,13 @@ def _set_dag_run_terminal_state( :param ti_state: state to set on running task instances :param commit: commit Dag and tasks to be altered to the database :param session: database session - :return: If commit is true, list of tasks that have been updated, - otherwise list of tasks that will be updated + :return: ``(all_updated_tis, running_tis)`` where ``all_updated_tis`` is the + combined list of pending (now SKIPPED) and running (now ``ti_state``) TIs, + and ``running_tis`` is only the TIs that were in an active state before + the transition (useful for firing terminal listener hooks). """ if not dag: - return [] + return [], [] if not run_id: raise ValueError(f"Invalid dag_run_id: {run_id}") @@ -302,13 +304,14 @@ def _set_running_task(task: Operator) -> Operator: if not any(dag.task_dict[ti.task_id].is_teardown for ti in (running_tis + pending_tis)): _set_dag_run_state(dag.dag_id, run_id, run_state, session) - return pending_normal_tis + set_state( + all_updated = pending_normal_tis + set_state( tasks=running_tasks, run_id=run_id, state=ti_state, commit=commit, session=session, ) + return all_updated, running_tis @provide_session @@ -318,7 +321,7 @@ def set_dag_run_state_to_success( run_id: str | None = None, commit: bool = False, session: SASession = NEW_SESSION, -) -> list[TaskInstance]: +) -> tuple[list[TaskInstance], list[TaskInstance]]: """ Set the dag run's state to success. @@ -328,8 +331,13 @@ def set_dag_run_state_to_success( :param run_id: the run_id to start looking from :param commit: commit Dag and tasks to be altered to the database :param session: database session - :return: If commit is true, list of tasks that have been updated, - otherwise list of tasks that will be updated + :return: A tuple ``(all_updated_tis, running_tis)``. + ``all_updated_tis`` is the combined list of task instances that were updated: + previously-running ones (now SUCCESS) and non-finished pending ones (now SKIPPED). + If ``commit`` is ``False``, this is the list of task instances *that would be* updated. + ``running_tis`` contains only the task instances that were in an active running state + (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) before the transition — + useful for firing terminal listener hooks. :raises: ValueError if dag or logical_date is invalid """ return _set_dag_run_terminal_state( @@ -349,7 +357,7 @@ def set_dag_run_state_to_failed( run_id: str | None = None, commit: bool = False, session: SASession = NEW_SESSION, -) -> list[TaskInstance]: +) -> tuple[list[TaskInstance], list[TaskInstance]]: """ Set the dag run's state to failed. @@ -359,8 +367,13 @@ def set_dag_run_state_to_failed( :param run_id: the Dag run_id to start looking from :param commit: commit Dag and tasks to be altered to the database :param session: database session - :return: If commit is true, list of tasks that have been updated, - otherwise list of tasks that will be updated + :return: A tuple ``(all_updated_tis, running_tis)``. + ``all_updated_tis`` is the combined list of task instances that were updated: + previously-running ones (now FAILED) and non-finished pending ones (now SKIPPED). + If ``commit`` is ``False``, this is the list of task instances *that would be* updated. + ``running_tis`` contains only the task instances that were in an active running state + (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) before the transition — + useful for firing terminal listener hooks. """ return _set_dag_run_terminal_state( dag=dag, diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py index 186cd3671c66b..ee0346fe56367 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py @@ -57,6 +57,7 @@ ) from airflow.api_fastapi.core_api.datamodels.task_instances import NewTaskResponse from airflow.api_fastapi.core_api.services.public.common import BulkService +from airflow.api_fastapi.core_api.services.public.task_instances import _emit_state_listener_hooks from airflow.listeners.listener import get_listener_manager from airflow.models.dagrun import DagRun, clear_partition_runs from airflow.models.taskinstance import TaskInstance @@ -193,7 +194,13 @@ def patch_dag_run_state( ) -> None: """Set a Dag Run's state (success/queued/failed), firing the matching listener hooks.""" if state == DagRunMutableStates.SUCCESS: - set_dag_run_state_to_success(dag=dag, run_id=dag_run.run_id, commit=True, session=session) + _, running_tis = set_dag_run_state_to_success( + dag=dag, run_id=dag_run.run_id, commit=True, session=session + ) + try: + _emit_state_listener_hooks(running_tis, TaskInstanceState.SUCCESS) + except Exception: + log.exception("error calling listener") try: get_listener_manager().hook.on_dag_run_success(dag_run=dag_run, msg="") except Exception: @@ -204,7 +211,13 @@ def patch_dag_run_state( # Not notifying on queued - only notifying on RUNNING, which happens in the scheduler. set_dag_run_state_to_queued(dag=dag, run_id=dag_run.run_id, commit=True, session=session) elif state == DagRunMutableStates.FAILED: - set_dag_run_state_to_failed(dag=dag, run_id=dag_run.run_id, commit=True, session=session) + _, running_tis = set_dag_run_state_to_failed( + dag=dag, run_id=dag_run.run_id, commit=True, session=session + ) + try: + _emit_state_listener_hooks(running_tis, TaskInstanceState.FAILED) + except Exception: + log.exception("error calling listener") try: get_listener_manager().hook.on_dag_run_failed(dag_run=dag_run, msg="") except Exception: 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 e1f5695ada8ce..b16582ea58e1b 100644 --- a/airflow-core/tests/unit/api/common/test_mark_tasks.py +++ b/airflow-core/tests/unit/api/common/test_mark_tasks.py @@ -50,9 +50,10 @@ def test_set_dag_run_state_to_failed(dag_maker: DagMaker[SerializedDAG]): ti.set_state(TaskInstanceState.RUNNING) dag_maker.session.flush() - updated_tis: list[TaskInstance] = set_dag_run_state_to_failed( + result: tuple[list[TaskInstance], list[TaskInstance]] = set_dag_run_state_to_failed( dag=dag, run_id=dr.run_id, commit=True, session=dag_maker.session ) + updated_tis, _ = result assert len(updated_tis) == 2 task_dict = {ti.task_id: ti for ti in updated_tis} assert task_dict["running"].state == TaskInstanceState.FAILED @@ -82,9 +83,10 @@ def test_set_dag_run_state_to_success_unfinished_teardown( dag_maker.session.flush() assert dr.state == DagRunState.RUNNING - updated_tis: list[TaskInstance] = set_dag_run_state_to_success( + result: tuple[list[TaskInstance], list[TaskInstance]] = set_dag_run_state_to_success( dag=dag, run_id=dr.run_id, commit=True, session=dag_maker.session ) + updated_tis, _ = result run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id)) assert run is not None assert run.state != DagRunState.SUCCESS @@ -111,9 +113,10 @@ def test_set_dag_run_state_to_success_keeps_finished_task_states( dag_maker.session.flush() dr.set_state(DagRunState.FAILED) - updated_tis: list[TaskInstance] = set_dag_run_state_to_success( + result: tuple[list[TaskInstance], list[TaskInstance]] = set_dag_run_state_to_success( dag=dag, run_id=dr.run_id, commit=True, session=dag_maker.session ) + updated_tis, _ = result run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id)) assert run is not None assert run.state == DagRunState.SUCCESS diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py index 3b1422c79c653..1dc5c5e3bdafb 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py @@ -48,7 +48,7 @@ from airflow.timetables.simple import PartitionedAssetTimetable, PartitionedAtRuntime from airflow.timetables.trigger import CronPartitionTimetable from airflow.utils.session import provide_session -from airflow.utils.state import DagRunState, State +from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType from tests_common.test_utils.api_fastapi import _check_dag_run_note, _check_last_log @@ -1609,6 +1609,70 @@ def test_patch_dag_run_notifies_listeners(self, test_client, state, listener_sta assert response.status_code == 200 assert listener.state == listener_state + @pytest.mark.parametrize( + ("dag_run_state", "expected_ti_listener_state"), + [ + ("success", TaskInstanceState.SUCCESS), + ("failed", TaskInstanceState.FAILED), + ], + ) + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_patch_dag_run_notifies_ti_listeners_for_running_tasks( + self, + test_client, + dag_maker, + session, + listener_manager, + dag_run_state, + expected_ti_listener_state, + ): + with dag_maker(dag_id="test_ti_listeners", schedule=None, serialized=True): + EmptyOperator(task_id="t1") + + dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) + ti = dr.get_task_instance(task_id="t1") + ti.state = TaskInstanceState.RUNNING + dag_maker.sync_dagbag_to_db() + session.commit() + + listener = ClassBasedListener() + listener_manager(listener) + + response = test_client.patch( + f"/dags/test_ti_listeners/dagRuns/{dr.run_id}", json={"state": dag_run_state} + ) + assert response.status_code == 200 + assert expected_ti_listener_state in listener.state + + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_patch_dag_run_does_not_notify_ti_listeners_for_non_running_tasks( + self, + test_client, + dag_maker, + session, + listener_manager, + ): + with dag_maker(dag_id="test_ti_listeners_queued", schedule=None, serialized=True): + EmptyOperator(task_id="t1") + + dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) + ti = dr.get_task_instance(task_id="t1") + ti.state = TaskInstanceState.QUEUED + dag_maker.sync_dagbag_to_db() + session.commit() + + listener = ClassBasedListener() + listener_manager(listener) + + response = test_client.patch( + f"/dags/test_ti_listeners_queued/dagRuns/{dr.run_id}", json={"state": "success"} + ) + assert response.status_code == 200 + # Only the dagrun-level hook should have fired; no TI hooks for a non-running task. + # DagRunState.SUCCESS and TaskInstanceState.SUCCESS share the string value 'success', so + # we check the exact type to distinguish them. + assert listener.state == [DagRunState.SUCCESS] + class TestDeleteDagRun: def test_delete_dag_run(self, test_client, session): From 8de2bf3faba40d588471a4fb54ec1458736a2c62 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 15 Jul 2026 15:45:51 +0200 Subject: [PATCH 2/2] Address PR review --- .../src/airflow/api/common/mark_tasks.py | 33 +++++--- .../core_api/services/public/dag_run.py | 14 +--- .../core_api/routes/public/test_dag_run.py | 81 ++++++++++++++++--- 3 files changed, 92 insertions(+), 36 deletions(-) diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py b/airflow-core/src/airflow/api/common/mark_tasks.py index 1e81a6b174f1d..9a07e1c9b6304 100644 --- a/airflow-core/src/airflow/api/common/mark_tasks.py +++ b/airflow-core/src/airflow/api/common/mark_tasks.py @@ -237,10 +237,12 @@ def _set_dag_run_terminal_state( :param ti_state: state to set on running task instances :param commit: commit Dag and tasks to be altered to the database :param session: database session - :return: ``(all_updated_tis, running_tis)`` where ``all_updated_tis`` is the + :return: ``(all_updated_tis, killed_tis)`` where ``all_updated_tis`` is the combined list of pending (now SKIPPED) and running (now ``ti_state``) TIs, - and ``running_tis`` is only the TIs that were in an active state before - the transition (useful for firing terminal listener hooks). + and ``killed_tis`` contains only the non-teardown TIs that were in an active + running state and were forcefully terminated (teardown TIs are intentionally + left running so they can finish their own cleanup, and must not receive a + terminal listener event here). """ if not dag: return [], [] @@ -266,9 +268,10 @@ def _set_dag_run_terminal_state( ) ).all() ) - # Do not kill teardown tasks task_ids_of_running_tis = {ti.task_id for ti in running_tis if not dag.task_dict[ti.task_id].is_teardown} + # Keep task instances that were killed with no cleanup capability (all running minus teardown ones). + killed_tis = [ti for ti in running_tis if ti.task_id in task_ids_of_running_tis] def _set_running_task(task: Operator) -> Operator: task.dag = dag @@ -311,7 +314,7 @@ def _set_running_task(task: Operator) -> Operator: commit=commit, session=session, ) - return all_updated, running_tis + return all_updated, killed_tis @provide_session @@ -331,13 +334,15 @@ def set_dag_run_state_to_success( :param run_id: the run_id to start looking from :param commit: commit Dag and tasks to be altered to the database :param session: database session - :return: A tuple ``(all_updated_tis, running_tis)``. + :return: A tuple ``(all_updated_tis, killed_tis)``. ``all_updated_tis`` is the combined list of task instances that were updated: previously-running ones (now SUCCESS) and non-finished pending ones (now SKIPPED). If ``commit`` is ``False``, this is the list of task instances *that would be* updated. - ``running_tis`` contains only the task instances that were in an active running state - (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) before the transition — - useful for firing terminal listener hooks. + ``killed_tis`` contains only the non-teardown task instances that were in an active + running state (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) and were + forcefully terminated — teardown tasks are intentionally excluded because they are + left running to finish their own cleanup. Use ``killed_tis`` for firing terminal + listener hooks. :raises: ValueError if dag or logical_date is invalid """ return _set_dag_run_terminal_state( @@ -367,13 +372,15 @@ def set_dag_run_state_to_failed( :param run_id: the Dag run_id to start looking from :param commit: commit Dag and tasks to be altered to the database :param session: database session - :return: A tuple ``(all_updated_tis, running_tis)``. + :return: A tuple ``(all_updated_tis, killed_tis)``. ``all_updated_tis`` is the combined list of task instances that were updated: previously-running ones (now FAILED) and non-finished pending ones (now SKIPPED). If ``commit`` is ``False``, this is the list of task instances *that would be* updated. - ``running_tis`` contains only the task instances that were in an active running state - (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) before the transition — - useful for firing terminal listener hooks. + ``killed_tis`` contains only the non-teardown task instances that were in an active + running state (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) and were + forcefully terminated — teardown tasks are intentionally excluded because they are + left running to finish their own cleanup. Use ``killed_tis`` for firing terminal + listener hooks. """ return _set_dag_run_terminal_state( dag=dag, diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py index ee0346fe56367..8044ec9b546a3 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py @@ -194,13 +194,10 @@ def patch_dag_run_state( ) -> None: """Set a Dag Run's state (success/queued/failed), firing the matching listener hooks.""" if state == DagRunMutableStates.SUCCESS: - _, running_tis = set_dag_run_state_to_success( + _, killed_tis = set_dag_run_state_to_success( dag=dag, run_id=dag_run.run_id, commit=True, session=session ) - try: - _emit_state_listener_hooks(running_tis, TaskInstanceState.SUCCESS) - except Exception: - log.exception("error calling listener") + _emit_state_listener_hooks(killed_tis, TaskInstanceState.SUCCESS) try: get_listener_manager().hook.on_dag_run_success(dag_run=dag_run, msg="") except Exception: @@ -211,13 +208,10 @@ def patch_dag_run_state( # Not notifying on queued - only notifying on RUNNING, which happens in the scheduler. set_dag_run_state_to_queued(dag=dag, run_id=dag_run.run_id, commit=True, session=session) elif state == DagRunMutableStates.FAILED: - _, running_tis = set_dag_run_state_to_failed( + _, killed_tis = set_dag_run_state_to_failed( dag=dag, run_id=dag_run.run_id, commit=True, session=session ) - try: - _emit_state_listener_hooks(running_tis, TaskInstanceState.FAILED) - except Exception: - log.exception("error calling listener") + _emit_state_listener_hooks(killed_tis, TaskInstanceState.FAILED) try: get_listener_manager().hook.on_dag_run_failed(dag_run=dag_run, msg="") except Exception: diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py index 1dc5c5e3bdafb..a1dbd79db090e 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py @@ -1594,23 +1594,25 @@ def test_patch_dag_run_bad_request(self, test_client): assert body["detail"][0]["msg"] == "Input should be 'queued', 'success' or 'failed'" @pytest.mark.parametrize( - ("state", "listener_state"), + ("state", "expected_dagrun_state"), [ - ("queued", []), - ("success", [DagRunState.SUCCESS]), - ("failed", [DagRunState.FAILED]), + ("queued", None), + ("success", DagRunState.SUCCESS), + ("failed", DagRunState.FAILED), ], ) @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") - def test_patch_dag_run_notifies_listeners(self, test_client, state, listener_state, listener_manager): + def test_patch_dag_run_notifies_listeners( + self, test_client, state, expected_dagrun_state, listener_manager + ): listener = ClassBasedListener() listener_manager(listener) response = test_client.patch(f"/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}", json={"state": state}) assert response.status_code == 200 - assert listener.state == listener_state + assert listener.state == ([expected_dagrun_state] if expected_dagrun_state else []) @pytest.mark.parametrize( - ("dag_run_state", "expected_ti_listener_state"), + ("dag_run_state", "expected_ti_state"), [ ("success", TaskInstanceState.SUCCESS), ("failed", TaskInstanceState.FAILED), @@ -1624,13 +1626,19 @@ def test_patch_dag_run_notifies_ti_listeners_for_running_tasks( session, listener_manager, dag_run_state, - expected_ti_listener_state, + expected_ti_state, ): with dag_maker(dag_id="test_ti_listeners", schedule=None, serialized=True): EmptyOperator(task_id="t1") dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) - ti = dr.get_task_instance(task_id="t1") + ti = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == dr.dag_id, + TaskInstance.run_id == dr.run_id, + TaskInstance.task_id == "t1", + ) + ) ti.state = TaskInstanceState.RUNNING dag_maker.sync_dagbag_to_db() session.commit() @@ -1642,7 +1650,10 @@ def test_patch_dag_run_notifies_ti_listeners_for_running_tasks( f"/dags/test_ti_listeners/dagRuns/{dr.run_id}", json={"state": dag_run_state} ) assert response.status_code == 200 - assert expected_ti_listener_state in listener.state + # Use `is` for type-safe comparison between TaskInstanceState and DagRunState. + assert listener.state[0] is expected_ti_state + assert listener.state[1] is DagRunState(dag_run_state) + assert len(listener.state) == 2 @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") def test_patch_dag_run_does_not_notify_ti_listeners_for_non_running_tasks( @@ -1656,7 +1667,13 @@ def test_patch_dag_run_does_not_notify_ti_listeners_for_non_running_tasks( EmptyOperator(task_id="t1") dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) - ti = dr.get_task_instance(task_id="t1") + ti = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == dr.dag_id, + TaskInstance.run_id == dr.run_id, + TaskInstance.task_id == "t1", + ) + ) ti.state = TaskInstanceState.QUEUED dag_maker.sync_dagbag_to_db() session.commit() @@ -1669,10 +1686,48 @@ def test_patch_dag_run_does_not_notify_ti_listeners_for_non_running_tasks( ) assert response.status_code == 200 # Only the dagrun-level hook should have fired; no TI hooks for a non-running task. - # DagRunState.SUCCESS and TaskInstanceState.SUCCESS share the string value 'success', so - # we check the exact type to distinguish them. + # The list length check distinguishes "only dagrun fired" from "both fired". assert listener.state == [DagRunState.SUCCESS] + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_patch_dag_run_does_not_notify_ti_listeners_for_running_teardown_tasks( + self, + test_client, + dag_maker, + session, + listener_manager, + ): + with dag_maker(dag_id="test_ti_listeners_teardown", schedule=None, serialized=True): + normal = EmptyOperator(task_id="normal") + teardown = EmptyOperator(task_id="teardown").as_teardown(setups=normal) + normal >> teardown + + dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) + for task_id in ("normal", "teardown"): + ti = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == dr.dag_id, + TaskInstance.run_id == dr.run_id, + TaskInstance.task_id == task_id, + ) + ) + ti.state = TaskInstanceState.RUNNING + dag_maker.sync_dagbag_to_db() + session.commit() + + listener = ClassBasedListener() + listener_manager(listener) + + response = test_client.patch( + f"/dags/test_ti_listeners_teardown/dagRuns/{dr.run_id}", json={"state": "success"} + ) + assert response.status_code == 200 + # Normal task was killed — its TI listener fires. Teardown task is intentionally skipped. + # Use `is` for type-safe comparison between TaskInstanceState and DagRunState. + assert len(listener.state) == 2 + assert listener.state[0] is TaskInstanceState.SUCCESS + assert listener.state[1] is DagRunState.SUCCESS + class TestDeleteDagRun: def test_delete_dag_run(self, test_client, session):