From eb16ce3cc17f5da19c2489d3c1c4419b17c23e2c Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 22 Jul 2026 17:53:35 +0200 Subject: [PATCH] fix: Patch ti note before state so that listeners can see it --- .../core_api/routes/public/task_instances.py | 73 +++++++++---------- .../services/public/task_instances.py | 32 ++++---- .../core_api/routes/public/test_dag_run.py | 2 +- .../routes/public/test_task_instances.py | 61 ++++++++++++++++ .../tests/unit/listeners/class_listener.py | 11 ++- 5 files changed, 122 insertions(+), 57 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py index f72cd3801c65c..7305e4e89be57 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -1013,6 +1013,14 @@ def patch_task_group_instances( ) response_tis = tis + # Apply "note" before "state" so listeners fired inside _patch_task_group_state() see the updated note. + if "note" in data: + _patch_task_instance_note( + task_instance_body=body, + tis=response_tis, + user=user, + update_mask=update_mask, + ) if "new_state" in data: response_tis = _patch_task_group_state( group_id=group_id, @@ -1022,13 +1030,6 @@ def patch_task_group_instances( data=data, session=session, ) - if "note" in data: - _patch_task_instance_note( - task_instance_body=body, - tis=response_tis, - user=user, - update_mask=update_mask, - ) response_tis = _reload_tis_with_rendered_fields(response_tis, session) @@ -1209,36 +1210,34 @@ def patch_task_instance( dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask ) - for key, _ in data.items(): - if key == "new_state": - # Create BulkTaskInstanceBody object with map_index field - bulk_ti_body = BulkTaskInstanceBody( - task_id=task_id, - map_index=map_index, - new_state=body.new_state, - note=body.note, - include_upstream=body.include_upstream, - include_downstream=body.include_downstream, - include_future=body.include_future, - include_past=body.include_past, - ) - - _patch_task_instance_state( - task_id=task_id, - dag_run_id=dag_run_id, - dag=dag, - task_instance_body=bulk_ti_body, - data=data, - session=session, - ) - - elif key == "note": - _patch_task_instance_note( - task_instance_body=body, - tis=tis, - user=user, - update_mask=update_mask, - ) + # Apply "note" before "state" so listeners fired inside _patch_task_instance_state() see the updated note. + if "note" in data: + _patch_task_instance_note( + task_instance_body=body, + tis=tis, + user=user, + update_mask=update_mask, + ) + if "new_state" in data: + # Create BulkTaskInstanceBody object with map_index field + bulk_ti_body = BulkTaskInstanceBody( + task_id=task_id, + map_index=map_index, + new_state=body.new_state, + note=body.note, + include_upstream=body.include_upstream, + include_downstream=body.include_downstream, + include_future=body.include_future, + include_past=body.include_past, + ) + _patch_task_instance_state( + task_id=task_id, + dag_run_id=dag_run_id, + dag=dag, + task_instance_body=bulk_ti_body, + data=data, + session=session, + ) return TaskInstanceCollectionResponse( task_instances=[ diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py index 0a1885f0e26fe..b00ba4effdbf6 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py @@ -468,22 +468,22 @@ def _perform_update( update_mask=update_mask, ) - for key, _ in data.items(): - if key == "new_state": - _patch_task_instance_state( - task_id=task_id, - dag_run_id=dag_run_id, - dag=dag, - task_instance_body=entity, - session=self.session, - data=data, - ) - elif key == "note": - _patch_task_instance_note( - task_instance_body=entity, - tis=tis, - user=self.user, - ) + # Apply "note" before "state" so listeners fired inside _patch_task_instance_state() see the updated note. + if "note" in data: + _patch_task_instance_note( + task_instance_body=entity, + tis=tis, + user=self.user, + ) + if "new_state" in data: + _patch_task_instance_state( + task_id=task_id, + dag_run_id=dag_run_id, + dag=dag, + task_instance_body=entity, + session=self.session, + data=data, + ) results.success.append(f"{dag_id}.{dag_run_id}.{task_id}[{map_index}]") 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 fb804e4c4c57d..0ae15bb0752b7 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 @@ -1612,7 +1612,7 @@ def test_patch_dag_run_notifies_listeners( 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 + assert listener.dag_run_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( diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py index b1f8f0ed0daec..ae3a8388feacd 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py @@ -4674,6 +4674,22 @@ def test_patch_task_instance_notifies_listeners( assert response2.json()["state"] == state assert listener.state == listener_state + def test_patch_task_instance_listener_sees_note_when_note_and_state_both_patched( + self, test_client, session, listener_manager + ): + from unit.listeners.class_listener import ClassBasedListener + + self.create_task_instances(session) + + listener = ClassBasedListener() + listener_manager(listener) + response = test_client.patch( + self.ENDPOINT_URL, + json={"new_state": "success", "note": "listener_note"}, + ) + assert response.status_code == 200 + assert listener.ti_note_at_listener == "listener_note" + @mock.patch("airflow.serialization.definitions.dag.SerializedDAG.set_task_instance_state") def test_should_call_mocked_api(self, mock_set_ti_state, test_client, session): self.create_task_instances(session) @@ -6976,6 +6992,35 @@ def test_should_respond_422(self, test_client): response = test_client.patch(self.ENDPOINT_URL, json={}) assert response.status_code == 422 + def test_bulk_update_listener_sees_note_when_note_and_state_both_patched( + self, test_client, session, listener_manager + ): + from unit.listeners.class_listener import ClassBasedListener + + self.create_task_instances(session, task_instances=[{"state": State.RUNNING}]) + + listener = ClassBasedListener() + listener_manager(listener) + response = test_client.patch( + self.ENDPOINT_URL, + json={ + "actions": [ + { + "action": "update", + "entities": [ + { + "task_id": self.TASK_ID, + "new_state": "success", + "note": "listener_note", + } + ], + } + ] + }, + ) + assert response.status_code == 200 + assert listener.ti_note_at_listener == "listener_note" + class TestPatchTaskGroup(TestTaskInstanceEndpoint): DAG_ID = "example_task_group" @@ -7421,6 +7466,22 @@ def test_patch_task_group_state_and_note(self, test_client, session): assert ti.state == TaskInstanceState.FAILED _check_task_instance_note(session, ti.id, {"content": note_value, "user_id": "test"}) + def test_patch_task_group_listener_sees_note_when_note_and_state_both_patched( + self, test_client, session, listener_manager + ): + from unit.listeners.class_listener import ClassBasedListener + + self.create_task_instances(session, dag_id=self.DAG_ID) + + listener = ClassBasedListener() + listener_manager(listener) + response = test_client.patch( + self.ENDPOINT_URL, + json={"new_state": "failed", "note": "listener_note"}, + ) + assert response.status_code == 200 + assert listener.ti_note_at_listener == "listener_note" + class TestPatchTaskGroupDryRun(TestTaskInstanceEndpoint): DAG_ID = "example_task_group" diff --git a/airflow-core/tests/unit/listeners/class_listener.py b/airflow-core/tests/unit/listeners/class_listener.py index 152f199cb806e..75d7d73d4bd64 100644 --- a/airflow-core/tests/unit/listeners/class_listener.py +++ b/airflow-core/tests/unit/listeners/class_listener.py @@ -27,8 +27,9 @@ def __init__(self): self.stopped_component = None self.state = [] self.dag_run_msg: str | None = None - self.dag_has_dag_attr: bool | None = None + self.dag_run_has_dag_attr: bool | None = None self.dag_run_note_at_listener: str | None = None + self.ti_note_at_listener: str | None = None @hookimpl def on_starting(self, component): @@ -43,18 +44,22 @@ def before_stopping(self, component): @hookimpl def on_task_instance_running(self, previous_state, task_instance): self.state.append(TaskInstanceState.RUNNING) + self.ti_note_at_listener = task_instance.note @hookimpl def on_task_instance_success(self, previous_state, task_instance): self.state.append(TaskInstanceState.SUCCESS) + self.ti_note_at_listener = task_instance.note @hookimpl def on_task_instance_failed(self, previous_state, task_instance, error: None | str | BaseException): self.state.append(TaskInstanceState.FAILED) + self.ti_note_at_listener = task_instance.note @hookimpl def on_task_instance_skipped(self, previous_state, task_instance): self.state.append(TaskInstanceState.SKIPPED) + self.ti_note_at_listener = task_instance.note @hookimpl def on_dag_run_running(self, dag_run, msg: str): @@ -64,14 +69,14 @@ def on_dag_run_running(self, dag_run, msg: str): 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_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_has_dag_attr = dag_run.dag is not None self.dag_run_note_at_listener = dag_run.note