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 b6d81b94ed9 Call listeners for running task instance when a Dag Run
state is manually set (#69874)
b6d81b94ed9 is described below
commit b6d81b94ed9977f374e7651aa1ee23cb4217d139
Author: Kacper Muda <[email protected]>
AuthorDate: Wed Jul 22 18:42:53 2026 +0200
Call listeners for running task instance when a Dag Run state is manually
set (#69874)
* Call listeners for running task instance when a Dag Run state is manually
set
* Address PR review
# Conflicts:
#
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
* Remove two comment lines.
* Add newsfragment
* fix after rebase
---
airflow-core/newsfragments/69874.bugfix.rst | 1 +
airflow-core/src/airflow/api/common/mark_tasks.py | 44 ++++++--
.../core_api/services/public/dag_run.py | 11 +-
.../tests/unit/api/common/test_mark_tasks.py | 9 +-
.../core_api/routes/public/test_dag_run.py | 123 ++++++++++++++++++++-
5 files changed, 167 insertions(+), 21 deletions(-)
diff --git a/airflow-core/newsfragments/69874.bugfix.rst
b/airflow-core/newsfragments/69874.bugfix.rst
new file mode 100644
index 00000000000..ad0d87336aa
--- /dev/null
+++ b/airflow-core/newsfragments/69874.bugfix.rst
@@ -0,0 +1 @@
+Listeners registered via ``on_task_instance_success`` /
``on_task_instance_failed`` are now called for non-teardown task instances that
were running when a Dag Run state is manually set to a terminal state (e.g. via
the API or UI).
diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py
b/airflow-core/src/airflow/api/common/mark_tasks.py
index 5274d8c706e..9a07e1c9b63 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_dag_run_terminal_state(
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 07181daa1c2..9f5e65501c0 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.dag_run import (
)
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:
if dag_run.dag is None:
dag_run.dag = dag
@@ -209,7 +213,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:
if dag_run.dag is None:
dag_run.dag = dag
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 e1f5695ada8..b16582ea58e 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 fb804e4c4c5..2c51cc50c4b 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.interval import
CronDataIntervalTimetable
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,7 +1594,7 @@ class TestPatchDagRun:
assert body["detail"][0]["msg"] == "Input should be 'queued',
'success' or 'failed'"
@pytest.mark.parametrize(
- ("state", "listener_state", "expected_msg"),
+ ("state", "expected_dagrun_state", "expected_msg"),
[
("queued", [], None),
("success", [DagRunState.SUCCESS], "Dag Run's state was manually
set to `success`."),
@@ -1603,13 +1603,13 @@ class TestPatchDagRun:
)
@pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
def test_patch_dag_run_notifies_listeners(
- self, test_client, state, listener_state, expected_msg,
listener_manager
+ self, test_client, state, expected_dagrun_state, expected_msg,
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_msg is not None:
assert listener.dag_run_msg == expected_msg
assert listener.dag_has_dag_attr is True
@@ -1627,6 +1627,121 @@ class TestPatchDagRun:
assert response.status_code == 200
assert listener.dag_run_note_at_listener == "listener_note"
+ @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
+ 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.
+ 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):
