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..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,13 +222,13 @@ 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(): - 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: 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():