diff --git a/airflow-core/docs/administration-and-deployment/listeners.rst b/airflow-core/docs/administration-and-deployment/listeners.rst index fec4edc94d2ee..2ae7899e26199 100644 --- a/airflow-core/docs/administration-and-deployment/listeners.rst +++ b/airflow-core/docs/administration-and-deployment/listeners.rst @@ -40,6 +40,8 @@ DagRun State Change Events -------------------------- DagRun state change events occur when a :class:`~airflow.models.dagrun.DagRun` changes state. +Beginning with Airflow 3, listeners are also notified whenever a state change is triggered through the API +(for ``on_dag_run_success`` and ``on_dag_run_failed``) e.g., when a DagRun is marked as success from the Airflow UI. - ``on_dag_run_running`` @@ -66,8 +68,12 @@ DagRun state change events occur when a :class:`~airflow.models.dagrun.DagRun` c TaskInstance State Change Events -------------------------------- -TaskInstance state change events occur when a :class:`~airflow.models.taskinstance.TaskInstance` changes state. +TaskInstance state change events occur when a :class:`~airflow.sdk.execution_time.task_runner.RuntimeTaskInstance` changes state. You can use these events to react to ``LocalTaskJob`` state changes. +Starting with Airflow 3, listeners are also notified when a state change is triggered through the API +(for ``on_task_instance_success`` and ``on_task_instance_failed``) e.g., when marking a task instance as success from the Airflow UI. +In such cases, the listener will receive a :class:`~airflow.models.taskinstance.TaskInstance` instance instead +of a :class:`~airflow.sdk.execution_time.task_runner.RuntimeTaskInstance` instance. - ``on_task_instance_running`` @@ -182,12 +188,14 @@ For example if you want to implement a listener that uses the ``error`` field in List of changes in the listener interfaces since 2.8.0 when they were introduced: -+-----------------+--------------------------------------------+-------------------------------------------------------------------------+ -| Airflow Version | Affected method | Change | -+=================+============================================+=========================================================================+ -| 2.10.0 | ``on_task_instance_failed`` | An error field added to the interface | -+-----------------+--------------------------------------------+-------------------------------------------------------------------------+ -| 3.0.0 | ``on_task_instance_running``, | ``session`` argument removed from task instance listeners, | -| | ``on_task_instance_success``, | ``task_instance`` object is now an instance of ``RuntimeTaskInstance`` | -| | ``on_task_instance_failed`` | | -+-----------------+--------------------------------------------+-------------------------------------------------------------------------+ ++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+ +| Airflow Version | Affected method | Change | ++=================+============================================+===============================================================================================================================+ +| 2.10.0 | ``on_task_instance_failed`` | An error field added to the interface | ++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+ +| 3.0.0 | ``on_task_instance_running`` | ``session`` argument removed from task instance listeners, | +| | | ``task_instance`` object is now an instance of ``RuntimeTaskInstance`` | ++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+ +| 3.0.0 | ``on_task_instance_failed``, | ``session`` argument removed from task instance listeners, | +| | ``on_task_instance_success`` | ``task_instance`` object is now an instance of ``RuntimeTaskInstance`` when on worker and ``TaskInstance`` when on API server | ++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+ 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 6a3afa74b4f79..540e1c49a9100 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 @@ -19,6 +19,7 @@ from typing import Annotated, Literal, cast +import structlog from fastapi import Depends, HTTPException, Query, Request, status from fastapi.exceptions import RequestValidationError from pydantic import ValidationError @@ -64,6 +65,7 @@ from airflow.api_fastapi.core_api.security import GetUserDep, ReadableTIFilterDep, requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import TaskNotFound +from airflow.listeners.listener import get_listener_manager from airflow.models import Base, DagRun from airflow.models.dag import DAG from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances @@ -73,6 +75,8 @@ from airflow.utils.db import get_query_count from airflow.utils.state import DagRunState, TaskInstanceState +log = structlog.get_logger(__name__) + task_instances_router = AirflowRouter(tags=["Task Instance"], prefix="/dags/{dag_id}") task_instances_prefix = "/dagRuns/{dag_run_id}/taskInstances" @@ -882,6 +886,20 @@ def patch_task_instance( status.HTTP_409_CONFLICT, f"Task id {task_id} is already in {data['new_state']} state" ) ti = tis[0] if isinstance(tis, list) else tis + try: + if data["new_state"] == TaskInstanceState.SUCCESS: + get_listener_manager().hook.on_task_instance_success( + previous_state=None, task_instance=ti + ) + elif data["new_state"] == TaskInstanceState.FAILED: + get_listener_manager().hook.on_task_instance_failed( + previous_state=None, + task_instance=ti, + error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.", + ) + except Exception: + log.exception("error calling listener") + elif key == "note": if update_mask or body.note is not None: if ti.task_instance_note is None: diff --git a/airflow-core/src/airflow/example_dags/plugins/event_listener.py b/airflow-core/src/airflow/example_dags/plugins/event_listener.py index e5793a0d6b33f..4e8581dc4af9c 100644 --- a/airflow-core/src/airflow/example_dags/plugins/event_listener.py +++ b/airflow-core/src/airflow/example_dags/plugins/event_listener.py @@ -20,6 +20,7 @@ from typing import TYPE_CHECKING from airflow.listeners import hookimpl +from airflow.models.taskinstance import TaskInstance if TYPE_CHECKING: from airflow.models.dagrun import DagRun @@ -31,10 +32,10 @@ @hookimpl def on_task_instance_running(previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance): """ - This method is called when task state changes to RUNNING. - Through callback, parameters like previous_task_state, task_instance object can be accessed. - This will give more information about current task_instance that is running its dag_run, - task and dag information. + Called when task state changes to RUNNING. + + previous_task_state and task_instance object can be used to retrieve more information about current + task_instance that is running, its dag_run, task and dag information. """ print("Task instance is in running state") print(" Previous state of the Task instance:", previous_state) @@ -61,16 +62,27 @@ def on_task_instance_running(previous_state: TaskInstanceState, task_instance: R # [START howto_listen_ti_success_task] @hookimpl -def on_task_instance_success(previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance): +def on_task_instance_success( + previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance | TaskInstance +): """ - This method is called when task state changes to SUCCESS. - Through callback, parameters like previous_task_state, task_instance object can be accessed. - This will give more information about current task_instance that has succeeded its - dag_run, task and dag information. + Called when task state changes to SUCCESS. + + previous_task_state and task_instance object can be used to retrieve more information about current + task_instance that has succeeded, its dag_run, task and dag information. + + A RuntimeTaskInstance is provided in most cases, except when the task's state change is triggered + through the API. In that case, the TaskInstance available on the API server will be provided instead. """ print("Task instance in success state") print(" Previous state of the Task instance:", previous_state) + if isinstance(task_instance, TaskInstance): + print("Task instance's state was changed through the API.") + + print(f"Task operator:{task_instance.operator}") + return + context = task_instance.get_template_context() operator = context["task"] @@ -83,16 +95,29 @@ def on_task_instance_success(previous_state: TaskInstanceState, task_instance: R # [START howto_listen_ti_failure_task] @hookimpl def on_task_instance_failed( - previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance, error: None | str | BaseException + previous_state: TaskInstanceState, + task_instance: RuntimeTaskInstance | TaskInstance, + error: None | str | BaseException, ): """ - This method is called when task state changes to FAILED. - Through callback, parameters like previous_task_state, task_instance object can be accessed. - This will give more information about current task_instance that has failed its dag_run, - task and dag information. + Called when task state changes to FAILED. + + previous_task_state, task_instance object and error can be used to retrieve more information about current + task_instance that has failed, its dag_run, task and dag information. + + A RuntimeTaskInstance is provided in most cases, except when the task's state change is triggered + through the API. In that case, the TaskInstance available on the API server will be provided instead. """ print("Task instance in failure state") + if isinstance(task_instance, TaskInstance): + print("Task instance's state was changed through the API.") + + print(f"Task operator:{task_instance.operator}") + if error: + print(f"Failure caused by {error}") + return + context = task_instance.get_template_context() task = context["task"] diff --git a/airflow-core/src/airflow/listeners/spec/taskinstance.py b/airflow-core/src/airflow/listeners/spec/taskinstance.py index d66d6c83ce3de..7e4bd5ecd8779 100644 --- a/airflow-core/src/airflow/listeners/spec/taskinstance.py +++ b/airflow-core/src/airflow/listeners/spec/taskinstance.py @@ -22,6 +22,7 @@ from pluggy import HookspecMarker if TYPE_CHECKING: + from airflow.models.taskinstance import TaskInstance from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance from airflow.utils.state import TaskInstanceState @@ -34,14 +35,16 @@ def on_task_instance_running(previous_state: TaskInstanceState | None, task_inst @hookspec -def on_task_instance_success(previous_state: TaskInstanceState | None, task_instance: RuntimeTaskInstance): +def on_task_instance_success( + previous_state: TaskInstanceState | None, task_instance: RuntimeTaskInstance | TaskInstance +): """Execute when task state changes to SUCCESS. previous_state can be None.""" @hookspec def on_task_instance_failed( previous_state: TaskInstanceState | None, - task_instance: RuntimeTaskInstance, + task_instance: RuntimeTaskInstance | TaskInstance, error: None | str | BaseException, ): """Execute when task state changes to FAIL. previous_state can be None.""" 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 4678e6c30eaf1..14e000a32927d 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 @@ -30,6 +30,7 @@ from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.jobs.job import Job from airflow.jobs.triggerer_job_runner import TriggererJobRunner +from airflow.listeners.listener import get_listener_manager from airflow.models import DagRun, TaskInstance from airflow.models.baseoperator import BaseOperator from airflow.models.dagbag import DagBag @@ -3032,6 +3033,39 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint): TASK_ID = "print_the_context" RUN_ID = "TEST_DAG_RUN_ID" + @pytest.fixture(autouse=True) + def clean_listener_manager(self): + get_listener_manager().clear() + yield + get_listener_manager().clear() + + @pytest.mark.parametrize( + "state, listener_state", + [ + ("success", [TaskInstanceState.SUCCESS]), + ("failed", [TaskInstanceState.FAILED]), + ("skipped", []), + ], + ) + def test_patch_task_instance_notifies_listeners(self, test_client, session, state, listener_state): + from unit.listeners.class_listener import ClassBasedListener + + self.create_task_instances(session) + + listener = ClassBasedListener() + get_listener_manager().add_listener(listener) + test_client.patch( + self.ENDPOINT_URL, + json={ + "new_state": state, + }, + ) + + response2 = test_client.get(self.ENDPOINT_URL) + assert response2.status_code == 200 + assert response2.json()["state"] == state + assert listener.state == listener_state + @mock.patch("airflow.models.dag.DAG.set_task_instance_state") def test_should_call_mocked_api(self, mock_set_ti_state, test_client, session): self.create_task_instances(session)