diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py b/airflow-core/src/airflow/api/common/mark_tasks.py index 5274d8c706e04..9a07e1c9b6304 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,15 @@ 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, killed_tis)`` where ``all_updated_tis`` is the + combined list of pending (now SKIPPED) and running (now ``ti_state``) TIs, + 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 [] + return [], [] if not run_id: raise ValueError(f"Invalid dag_run_id: {run_id}") @@ -264,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 @@ -302,13 +307,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, killed_tis @provide_session @@ -318,7 +324,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 +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: 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, 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. + ``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( @@ -349,7 +362,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 +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: 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, 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. + ``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 186cd3671c66b..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 @@ -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,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: - set_dag_run_state_to_success(dag=dag, run_id=dag_run.run_id, commit=True, session=session) + _, killed_tis = set_dag_run_state_to_success( + dag=dag, run_id=dag_run.run_id, commit=True, session=session + ) + _emit_state_listener_hooks(killed_tis, TaskInstanceState.SUCCESS) try: get_listener_manager().hook.on_dag_run_success(dag_run=dag_run, msg="") except Exception: @@ -204,7 +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: - set_dag_run_state_to_failed(dag=dag, run_id=dag_run.run_id, commit=True, session=session) + _, killed_tis = set_dag_run_state_to_failed( + dag=dag, run_id=dag_run.run_id, commit=True, session=session + ) + _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/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..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 @@ -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 @@ -1594,20 +1594,139 @@ 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_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_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 = 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() + + 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 + # 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( + 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 = 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() + + 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. + # 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: