From c5c0a901be268f7f69133633ea05275b96fa977d Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Tue, 14 Jul 2026 15:17:39 +0200 Subject: [PATCH 1/2] Populate dag and note before firing dag run state-change listeners --- .../core_api/routes/public/dag_run.py | 4 ++- .../core_api/services/public/dag_run.py | 14 ++++++++-- .../core_api/routes/public/test_dag_run.py | 28 +++++++++++++++---- .../tests/unit/listeners/class_listener.py | 9 ++++++ 4 files changed, 47 insertions(+), 8 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py index 81d47539229ce..5295b1e72550b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -222,7 +222,9 @@ def patch_dag_run( data = patch_body.model_dump(include=fields_to_update, by_alias=True) - for attr_name, attr_value_raw in data.items(): + # Sort keys so that the "note" update happens before "state" update, so that the listeners called inside + # `patch_dag_run_state()` already receive updated DagRun object with note. + for attr_name, attr_value_raw in sorted(data.items()): if attr_name == "state" and patch_body.state is not None: patch_dag_run_state(dag=dag, dag_run=dag_run, state=patch_body.state, session=session) elif attr_name == "note": 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..07181daa1c214 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 @@ -195,7 +195,12 @@ def patch_dag_run_state( if state == DagRunMutableStates.SUCCESS: set_dag_run_state_to_success(dag=dag, run_id=dag_run.run_id, commit=True, session=session) try: - get_listener_manager().hook.on_dag_run_success(dag_run=dag_run, msg="") + if dag_run.dag is None: + dag_run.dag = dag + get_listener_manager().hook.on_dag_run_success( + dag_run=dag_run, + msg=f"Dag Run's state was manually set to `{DagRunMutableStates.SUCCESS.value}`.", + ) except Exception: log.exception("error calling listener") elif state == DagRunMutableStates.QUEUED: @@ -206,7 +211,12 @@ def patch_dag_run_state( elif state == DagRunMutableStates.FAILED: set_dag_run_state_to_failed(dag=dag, run_id=dag_run.run_id, commit=True, session=session) try: - get_listener_manager().hook.on_dag_run_failed(dag_run=dag_run, msg="") + if dag_run.dag is None: + dag_run.dag = dag + get_listener_manager().hook.on_dag_run_failed( + dag_run=dag_run, + msg=f"Dag Run's state was manually set to `{DagRunMutableStates.FAILED.value}`.", + ) except Exception: log.exception("error calling listener") 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..fb804e4c4c57d 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,20 +1594,38 @@ 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", "listener_state", "expected_msg"), [ - ("queued", []), - ("success", [DagRunState.SUCCESS]), - ("failed", [DagRunState.FAILED]), + ("queued", [], None), + ("success", [DagRunState.SUCCESS], "Dag Run's state was manually set to `success`."), + ("failed", [DagRunState.FAILED], "Dag Run's state was manually set to `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, listener_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 + if expected_msg is not None: + assert listener.dag_run_msg == expected_msg + assert listener.dag_has_dag_attr is True + + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_patch_dag_run_listener_sees_note_when_note_and_state_both_patched( + self, test_client, listener_manager + ): + listener = ClassBasedListener() + listener_manager(listener) + response = test_client.patch( + f"/dags/{DAG1_ID}/dagRuns/{DAG1_RUN2_ID}", + json={"state": "success", "note": "listener_note"}, + ) + assert response.status_code == 200 + assert listener.dag_run_note_at_listener == "listener_note" class TestDeleteDagRun: diff --git a/airflow-core/tests/unit/listeners/class_listener.py b/airflow-core/tests/unit/listeners/class_listener.py index 4b9ef3a9f26f1..152f199cb806e 100644 --- a/airflow-core/tests/unit/listeners/class_listener.py +++ b/airflow-core/tests/unit/listeners/class_listener.py @@ -26,6 +26,9 @@ def __init__(self): self.started_component = None self.stopped_component = None self.state = [] + self.dag_run_msg: str | None = None + self.dag_has_dag_attr: bool | None = None + self.dag_run_note_at_listener: str | None = None @hookimpl def on_starting(self, component): @@ -60,10 +63,16 @@ def on_dag_run_running(self, dag_run, msg: str): @hookimpl def on_dag_run_success(self, dag_run, msg: str): self.state.append(DagRunState.SUCCESS) + self.dag_run_msg = msg + self.dag_has_dag_attr = dag_run.dag is not None + self.dag_run_note_at_listener = dag_run.note @hookimpl def on_dag_run_failed(self, dag_run, msg: str): self.state.append(DagRunState.FAILED) + self.dag_run_msg = msg + self.dag_has_dag_attr = dag_run.dag is not None + self.dag_run_note_at_listener = dag_run.note def clear(): From 1cdf0dc2e44851e85717648677e2c14d74ce1845 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 15 Jul 2026 13:38:52 +0200 Subject: [PATCH 2/2] Address PR review --- .../core_api/routes/public/dag_run.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py index 5295b1e72550b..c7adb60e5e059 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -222,15 +222,13 @@ def patch_dag_run( data = patch_body.model_dump(include=fields_to_update, by_alias=True) - # Sort keys so that the "note" update happens before "state" update, so that the listeners called inside - # `patch_dag_run_state()` already receive updated DagRun object with note. - for attr_name, attr_value_raw in sorted(data.items()): - if attr_name == "state" and patch_body.state is not None: - patch_dag_run_state(dag=dag, dag_run=dag_run, state=patch_body.state, session=session) - elif attr_name == "note": - updated_dag_run = session.get(DagRun, dag_run.id) - if updated_dag_run is not None: - patch_dag_run_note(dag_run=updated_dag_run, note=attr_value_raw, user=user) + # Apply "note" before "state" so listeners fired inside patch_dag_run_state() see the updated note. + if "note" in data: + updated_dag_run = session.get(DagRun, dag_run.id) + if updated_dag_run is not None: + patch_dag_run_note(dag_run=updated_dag_run, note=data["note"], user=user) + if "state" in data and patch_body.state is not None: + patch_dag_run_state(dag=dag, dag_run=dag_run, state=patch_body.state, session=session) final_dag_run = session.get(DagRun, dag_run.id) if not final_dag_run: