From 3f2a4e6a910c615d45861df50e5ee0e87291f922 Mon Sep 17 00:00:00 2001 From: bingqin2 Date: Sun, 13 Sep 2026 09:56:53 -0500 Subject: [PATCH 1/3] Check that external tasks and task groups exist on Airflow 3 in ExternalTaskSensor --- .../standard/sensors/external_task.py | 52 +++++- .../standard/triggers/external_task.py | 38 ++++ .../providers/standard/utils/sensor_helper.py | 78 ++++++++ .../sensors/test_external_task_sensor.py | 176 ++++++++++++++++++ .../standard/triggers/test_external_task.py | 100 ++++++++++ 5 files changed, 442 insertions(+), 2 deletions(-) diff --git a/providers/standard/src/airflow/providers/standard/sensors/external_task.py b/providers/standard/src/airflow/providers/standard/sensors/external_task.py index 034e86ca2612c..a521495ffb2a0 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py +++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py @@ -43,7 +43,11 @@ ) from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.triggers.external_task import WorkflowTrigger -from airflow.providers.standard.utils.sensor_helper import _get_count, _get_external_task_group_task_ids +from airflow.providers.standard.utils.sensor_helper import ( + _check_external_task_existence, + _get_count, + _get_external_task_group_task_ids, +) from airflow.providers.standard.version_compat import ( AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS, @@ -64,6 +68,7 @@ from sqlalchemy.orm import Session from airflow.providers.common.compat.sdk import Context, TaskInstanceKey + from airflow.sdk.types import RuntimeTaskInstanceProtocol class ExternalDagLink(BaseOperatorLink): @@ -171,6 +176,11 @@ class ExternalTaskSensor(BaseSensorOperator): external_task_id is not None) or check if the DAG to wait for exists (when external_task_id is None), and immediately cease waiting if the external task or DAG does not exist (default value: False). + On Airflow 3 a worker has no database access, so tasks and task groups are checked + through the execution API against each awaited Dag run once that run exists: a run's + task instances are created together with it from the run's own Dag version. Until the + run exists the sensor keeps waiting. Whether the Dag itself is registered is not checked + on Airflow 3. :param poke_interval: polling period in seconds to check for the status :param poll_interval: (DEPRECATED) use ``poke_interval`` instead :param deferrable: Run sensor in deferrable mode @@ -351,8 +361,9 @@ def poke(self, context: Context) -> bool: def _poke_af3(self, context: Context, dttm_filter: Sequence[datetime.datetime]) -> bool: from airflow.providers.standard.utils.sensor_helper import _get_count_by_matched_states - self._has_checked_existence = True ti = context["ti"] + if self.check_existence and not self._has_checked_existence: + self._check_for_existence_af3(ti, dttm_filter) def _get_count(states: list[str]) -> int: if self.external_task_ids: @@ -488,6 +499,7 @@ def execute(self, context: Context) -> None: logical_dates=list(dttm_filter), run_ids=None, execution_dates=None, + check_existence=self.check_existence, ), method_name="execute_complete", ) @@ -533,6 +545,15 @@ def execute_complete(self, context: Context, event: dict[str, typing.Any] | None if self.soft_fail: raise AirflowSkipException("External job has failed skipping.") raise ExternalDagFailedError("External job has failed.") + elif event["status"] == "not_found": + # A missing task or task group is a configuration error rather than a sensor failure, + # so like the poke path this is raised regardless of soft_fail. + message = event.get("message") or ( + f"The external tasks or task group awaited in Dag {self.external_dag_id} do not exist." + ) + if event.get("kind") == "task_group": + raise ExternalTaskGroupNotFoundError(message) + raise ExternalTaskNotFoundError(message) else: if self.soft_fail: raise AirflowSkipException("External job has failed skipping.") @@ -541,6 +562,33 @@ def execute_complete(self, context: Context, event: dict[str, typing.Any] | None "name of executed task and Dag." ) + def _check_for_existence_af3( + self, ti: RuntimeTaskInstanceProtocol, dttm_filter: Sequence[datetime.datetime] + ) -> None: + """ + Check that the awaited tasks or task group exist, through the execution API. + + A worker has no database access on Airflow 3, so unlike ``_check_for_existence`` this + relies on what the execution API exposes. A run's task instances are created together + with it from the run's own Dag version, so once an awaited run exists they answer + whether a task belongs to that run. While a run does not exist yet nothing can be + concluded, so the check is repeated on later pokes until every awaited run has been seen. + + :param ti: the task instance running this sensor, used to reach the execution API + :param dttm_filter: the logical dates of the awaited Dag runs + """ + if not self.external_task_ids and not self.external_task_group_id: + self._has_checked_existence = True + return + + self._has_checked_existence = _check_external_task_existence( + ti, + external_dag_id=self.external_dag_id, + external_task_ids=self.external_task_ids, + external_task_group_id=self.external_task_group_id, + logical_dates=list(dttm_filter), + ) + def _check_for_existence(self, session: Session) -> None: dag_to_wait = DagModel.get_current(self.external_dag_id, session=session) diff --git a/providers/standard/src/airflow/providers/standard/triggers/external_task.py b/providers/standard/src/airflow/providers/standard/triggers/external_task.py index 8a8e7f9db6bfe..0e64cdb652d0a 100644 --- a/providers/standard/src/airflow/providers/standard/triggers/external_task.py +++ b/providers/standard/src/airflow/providers/standard/triggers/external_task.py @@ -25,6 +25,7 @@ from sqlalchemy import func, select from airflow.models import DagRun +from airflow.providers.standard.exceptions import ExternalTaskGroupNotFoundError, ExternalTaskNotFoundError from airflow.providers.standard.utils.sensor_helper import _get_count from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS from airflow.triggers.base import BaseTrigger, TriggerEvent @@ -51,6 +52,8 @@ class WorkflowTrigger(BaseTrigger): :param poke_interval: The interval (in seconds) for poking the external tasks. :param soft_fail: If True, the trigger will not fail the entire dag on external task failure. :param logical_dates: A list of logical dates for the external dag. + :param check_existence: If True, verify that the external tasks or task group exist in each + awaited Dag run once that run exists, and fire a ``not_found`` event otherwise. Airflow 3 only. """ def __init__( @@ -66,6 +69,7 @@ def __init__( allowed_states: Collection[str] | None = None, poke_interval: float = 2.0, soft_fail: bool = False, + check_existence: bool = False, **kwargs, ): self.external_dag_id = external_dag_id @@ -79,6 +83,7 @@ def __init__( self.soft_fail = soft_fail self.execution_dates = execution_dates self.logical_dates = logical_dates + self.check_existence = check_existence super().__init__(**kwargs) def serialize(self) -> tuple[str, dict[str, Any]]: @@ -96,6 +101,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: if AIRFLOW_V_3_0_PLUS: data["run_ids"] = self.run_ids data["logical_dates"] = self.logical_dates + data["check_existence"] = self.check_existence else: data["execution_dates"] = self.execution_dates @@ -110,7 +116,25 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]: get_count_func = self._get_count run_id_or_dates = self.execution_dates or [] + # Tasks and task groups can only be verified against a Dag run that exists, so this is + # repeated at every poll until every awaited run has been seen. + existence_checked = not ( + AIRFLOW_V_3_0_PLUS + and self.check_existence + and (self.external_task_ids or self.external_task_group_id) + ) + while True: + if not existence_checked: + try: + existence_checked = await self._check_existence_af_3() + except ExternalTaskGroupNotFoundError as e: + yield TriggerEvent({"status": "not_found", "kind": "task_group", "message": str(e)}) + return + except ExternalTaskNotFoundError as e: + yield TriggerEvent({"status": "not_found", "kind": "task", "message": str(e)}) + return + if self.failed_states: failed_count = await get_count_func(self.failed_states) if failed_count > 0: @@ -130,6 +154,20 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]: self.log.info("Sleeping for %s seconds", self.poke_interval) await asyncio.sleep(self.poke_interval) + async def _check_existence_af_3(self) -> bool: + """Check the awaited tasks or task group against the runs that exist; True once every run was checked.""" + from airflow.providers.standard.utils.sensor_helper import _check_external_task_existence + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + return await sync_to_async(_check_external_task_existence)( + RuntimeTaskInstance, + external_dag_id=self.external_dag_id, + external_task_ids=self.external_task_ids, + external_task_group_id=self.external_task_group_id, + logical_dates=self.logical_dates, + run_ids=self.run_ids, + ) + async def _get_count_af_3(self, states: Collection[str] | None) -> int: from airflow.providers.standard.utils.sensor_helper import _get_count_by_matched_states from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance diff --git a/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py b/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py index f16e86bf34a5e..6571064e5c43f 100644 --- a/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py +++ b/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py @@ -17,6 +17,7 @@ from __future__ import annotations from collections.abc import Collection +from http import HTTPStatus from typing import TYPE_CHECKING, Any, cast from sqlalchemy import func, select, tuple_ @@ -27,6 +28,8 @@ from airflow.utils.session import NEW_SESSION, provide_session if TYPE_CHECKING: + from datetime import datetime + from sqlalchemy.orm import Session from sqlalchemy.sql import Select @@ -153,3 +156,78 @@ def _get_count_by_matched_states( count += 1 return count + + +def _check_external_task_existence( + api: Any, + *, + external_dag_id: str, + external_task_ids: Collection[str] | None, + external_task_group_id: str | None, + logical_dates: Collection[datetime] | None = None, + run_ids: Collection[str] | None = None, +) -> bool: + """ + Verify that the awaited tasks or task group exist in the awaited Dag runs, through the execution API. + + A Dag run's task instances are created together with the run, in the same transaction and + from the run's own Dag version, so once a run exists its task instances are the + version-accurate answer to whether a task belongs to it. Nothing can be concluded about a run + that does not exist yet, which is why the caller has to repeat the check until this function + returns True. + + :param api: an object exposing ``get_dr_count``, ``get_ti_count`` and ``get_task_states`` the + way ``RuntimeTaskInstance`` does: the running task instance, or the class itself. + :param external_dag_id: The ID of the external Dag. + :param external_task_ids: The task IDs that must exist in every awaited run. + :param external_task_group_id: The task group ID that must exist in every awaited run. + :param logical_dates: Logical dates identifying the awaited runs, used when ``run_ids`` is empty. + :param run_ids: Run IDs identifying the awaited runs. + :return: True once every awaited run exists and passed the check, False while at least one + awaited run does not exist yet. + :raises ExternalTaskNotFoundError: when an existing run has no task instance for one of the tasks. + :raises ExternalTaskGroupNotFoundError: when the Dag has no such task group, or an existing run + has no task instance for any task of the group. + """ + from airflow.providers.standard.exceptions import ( + ExternalTaskGroupNotFoundError, + ExternalTaskNotFoundError, + ) + from airflow.sdk.exceptions import AirflowRuntimeError + + awaited_runs: list[tuple[str, dict[str, list[Any]]]] + if run_ids: + awaited_runs = [(run_id, {"run_ids": [run_id]}) for run_id in run_ids] + else: + awaited_runs = [(dt.isoformat(), {"logical_dates": [dt]}) for dt in logical_dates or []] + + all_runs_checked = True + for run_label, run_filter in awaited_runs: + if api.get_dr_count(dag_id=external_dag_id, **run_filter) == 0: + all_runs_checked = False + continue + for task_id in external_task_ids or (): + if api.get_ti_count(dag_id=external_dag_id, task_ids=[task_id], **run_filter) == 0: + raise ExternalTaskNotFoundError( + f"The external task {task_id} in Dag {external_dag_id} does not exist for run {run_label}." + ) + + if external_task_group_id: + try: + run_id_task_state_map = api.get_task_states( + dag_id=external_dag_id, task_group_id=external_task_group_id, **run_filter + ) + except AirflowRuntimeError as e: + if (e.error.detail or {}).get("status_code") == HTTPStatus.NOT_FOUND: + raise ExternalTaskGroupNotFoundError( + f"The external task group '{external_task_group_id}' in Dag '{external_dag_id}' " + "does not exist." + ) from None + raise + if not any(run_id_task_state_map.values()): + raise ExternalTaskGroupNotFoundError( + f"The external task group '{external_task_group_id}' in Dag '{external_dag_id}' " + f"does not exist for run {run_label}." + ) + + return all_runs_checked diff --git a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py index 4ef032c92daac..7e77fa60be527 100644 --- a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py +++ b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py @@ -1187,6 +1187,19 @@ def func(dt, context): @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Different test for AF 2") @pytest.mark.usefixtures("testing_dag_bundle") +def _api_server_error(status_code: int): + """Build the error a task gets back when an execution API call fails.""" + from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType + from airflow.sdk.execution_time.comms import ErrorResponse + + return AirflowRuntimeError( + ErrorResponse( + error=ErrorType.API_SERVER_ERROR, + detail={"status_code": status_code, "message": f"Server returned {status_code}"}, + ) + ) + + class TestExternalTaskSensorV3: def setup_method(self): # Create a mock for TaskInstance with get_ti_count method @@ -1640,6 +1653,169 @@ def func(dt, context): context = {"logical_date": DEFAULT_DATE} assert op._handle_execution_date_fn(context) == DEFAULT_DATE + @pytest.mark.execution_timeout(10) + def test_check_existence_waits_for_run_then_checks_task_once(self, dag_maker): + """Nothing is concluded until the awaited run exists; once it does, the task is checked once.""" + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_id="test_task", + check_existence=True, + ) + + self.context["ti"].get_dr_count.side_effect = [0, 1] + # the existence probe (no ``states``) finds the task's instance; the state poll keeps waiting + self.context["ti"].get_ti_count.side_effect = lambda **kwargs: 0 if "states" in kwargs else 1 + + assert op.poke(self.context) is False # no run yet: nothing can be checked + assert op.poke(self.context) is False # run exists: task verified, still waiting for its state + assert op.poke(self.context) is False # already verified: no further run lookup + + assert self.context["ti"].get_dr_count.call_count == 2 + existence_probes = [ + c for c in self.context["ti"].get_ti_count.call_args_list if "states" not in c.kwargs + ] + assert existence_probes == [ + mock.call(dag_id="test_dag_parent", task_ids=["test_task"], logical_dates=[DEFAULT_DATE]), + ] + + @pytest.mark.execution_timeout(10) + def test_check_existence_task_not_found_in_run(self, dag_maker): + """A run that exists but has no task instance for the awaited task fails the sensor.""" + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_id="missing_task", + check_existence=True, + ) + + self.context["ti"].get_dr_count.return_value = 1 + self.context["ti"].get_ti_count.return_value = 0 + + with pytest.raises(ExternalTaskNotFoundError, match="missing_task"): + op.poke(self.context) + + self.context["ti"].get_ti_count.assert_called_once_with( + dag_id="test_dag_parent", task_ids=["missing_task"], logical_dates=[DEFAULT_DATE] + ) + + @pytest.mark.execution_timeout(10) + def test_check_existence_task_not_found_ignores_soft_fail(self, dag_maker): + """A missing task is a configuration error, so soft_fail does not turn it into a skip.""" + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_id="missing_task", + check_existence=True, + soft_fail=True, + ) + + self.context["ti"].get_dr_count.return_value = 1 + self.context["ti"].get_ti_count.return_value = 0 + + with pytest.raises(ExternalTaskNotFoundError, match="missing_task"): + op.execute(context=self.context) + + @pytest.mark.execution_timeout(10) + def test_check_existence_task_group_not_found_in_run(self, dag_maker): + """A run that exists without any task instance of the awaited group fails the sensor.""" + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_group_id="missing_group", + check_existence=True, + ) + + self.context["ti"].get_dr_count.return_value = 1 + self.context["ti"].get_task_states.return_value = {} + + with pytest.raises(ExternalTaskGroupNotFoundError, match="missing_group"): + op.poke(self.context) + + self.context["ti"].get_task_states.assert_called_once_with( + dag_id="test_dag_parent", task_group_id="missing_group", logical_dates=[DEFAULT_DATE] + ) + + @pytest.mark.execution_timeout(10) + def test_check_existence_task_group_unknown_to_dag(self, dag_maker): + """The execution API answers 404 for a task group the Dag does not define.""" + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_group_id="missing_group", + check_existence=True, + ) + + self.context["ti"].get_dr_count.return_value = 1 + self.context["ti"].get_task_states.side_effect = _api_server_error(404) + + with pytest.raises(ExternalTaskGroupNotFoundError, match="missing_group"): + op.poke(self.context) + + @pytest.mark.execution_timeout(10) + def test_check_existence_task_group_other_error_is_not_swallowed(self, dag_maker): + """Errors other than a missing task group keep their own type.""" + from airflow.sdk.exceptions import AirflowRuntimeError + + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_group_id="test_group", + check_existence=True, + ) + + self.context["ti"].get_dr_count.return_value = 1 + self.context["ti"].get_task_states.side_effect = _api_server_error(500) + + with pytest.raises(AirflowRuntimeError): + op.poke(self.context) + + @pytest.mark.execution_timeout(10) + def test_check_existence_deferrable_leaves_check_to_trigger(self, dag_maker): + """Nothing is checked before deferring; the trigger checks the tasks once their runs exist.""" + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_id="test_task", + deferrable=True, + check_existence=True, + ) + + with pytest.raises(TaskDeferred) as exc: + op.execute(context=self.context) + + assert exc.value.trigger.check_existence is True + self.context["ti"].get_dr_count.assert_not_called() + + @pytest.mark.parametrize( + ("kind", "expected_exception"), + [("task", ExternalTaskNotFoundError), ("task_group", ExternalTaskGroupNotFoundError)], + ) + def test_execute_complete_not_found_event(self, dag_maker, kind, expected_exception): + """A not_found event from the trigger raises the matching exception, even with soft_fail.""" + with dag_maker("test_dag_child"): + op = ExternalTaskSensor( + task_id="test_external_task_sensor_check", + external_dag_id="test_dag_parent", + external_task_id="test_task", + deferrable=True, + check_existence=True, + soft_fail=True, + ) + + with pytest.raises(expected_exception, match="does not exist"): + op.execute_complete( + context=self.context, + event={"status": "not_found", "kind": kind, "message": "The external thing does not exist."}, + ) + class TestExternalTaskAsyncSensor: TASK_ID = "external_task_sensor_check" diff --git a/providers/standard/tests/unit/standard/triggers/test_external_task.py b/providers/standard/tests/unit/standard/triggers/test_external_task.py index 09ba4ca6950f5..d8c0661d5ac8d 100644 --- a/providers/standard/tests/unit/standard/triggers/test_external_task.py +++ b/providers/standard/tests/unit/standard/triggers/test_external_task.py @@ -407,8 +407,108 @@ def test_serialization(self): "allowed_states": self.STATES, "poke_interval": 5, "soft_fail": False, + "check_existence": False, } + def test_serialization_check_existence(self): + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + external_task_ids=[self.TASK_ID], + check_existence=True, + ) + _, kwargs = trigger.serialize() + assert kwargs["check_existence"] is True + + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_dr_count") + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @pytest.mark.asyncio + async def test_task_workflow_trigger_task_not_found(self, mock_get_ti_count, mock_get_dr_count): + """A run that exists without the awaited task ends the wait with a not_found event.""" + mock_get_dr_count.return_value = 1 + mock_get_ti_count.return_value = 0 # the run has no task instance for the awaited task + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + external_task_ids=[self.TASK_ID], + allowed_states=self.STATES, + poke_interval=0.2, + check_existence=True, + ) + gen = trigger.run() + result = await gen.__anext__() + assert result.payload["status"] == "not_found" + assert result.payload["kind"] == "task" + assert self.TASK_ID in result.payload["message"] + mock_get_dr_count.assert_called_once_with(dag_id=self.DAG_ID, logical_dates=[self.LOGICAL_DATE]) + mock_get_ti_count.assert_any_call( + dag_id=self.DAG_ID, task_ids=[self.TASK_ID], logical_dates=[self.LOGICAL_DATE] + ) + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_dr_count") + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_task_states") + @pytest.mark.asyncio + async def test_task_workflow_trigger_task_group_not_found( + self, mock_get_task_states, mock_get_ti_count, mock_get_dr_count + ): + """A run that exists without any task of the awaited group ends the wait with a not_found event.""" + mock_get_dr_count.return_value = 1 + mock_get_task_states.return_value = {} + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + run_ids=[self.RUN_ID], + external_task_group_id="missing_group", + allowed_states=self.STATES, + poke_interval=0.2, + check_existence=True, + ) + gen = trigger.run() + result = await gen.__anext__() + assert result.payload["status"] == "not_found" + assert result.payload["kind"] == "task_group" + assert "missing_group" in result.payload["message"] + mock_get_task_states.assert_called_once_with( + dag_id=self.DAG_ID, task_group_id="missing_group", run_ids=[self.RUN_ID] + ) + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_dr_count") + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @mock.patch("asyncio.sleep") + @pytest.mark.asyncio + async def test_task_workflow_trigger_waits_for_run_before_checking( + self, mock_sleep, mock_get_ti_count, mock_get_dr_count + ): + """Nothing is concluded while the run is missing; once it exists the task is verified and the wait goes on.""" + mock_get_dr_count.side_effect = [0, 1] + + def ti_count(**kwargs): + if "states" in kwargs: + # the state poll: the task reaches an allowed state only after the run exists + return 1 if mock_get_dr_count.call_count == 2 else 0 + return 1 # existence probes: the run has task instances, including the awaited task + + mock_get_ti_count.side_effect = ti_count + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + external_task_ids=[self.TASK_ID], + allowed_states=self.STATES, + poke_interval=0.2, + check_existence=True, + ) + gen = trigger.run() + result = await gen.__anext__() + assert result.payload == {"status": "success"} + assert mock_get_dr_count.call_count == 2 + mock_sleep.assert_awaited_once() + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 2") class TestWorkflowTriggerAF2: From 2e86dd52af388389ffa353e204161bcfceff6918 Mon Sep 17 00:00:00 2001 From: bingqin2 Date: Sun, 13 Sep 2026 11:11:35 -0500 Subject: [PATCH 2/3] Keep the Airflow 3 sensor tests skipped on Airflow 2 --- .../tests/unit/standard/sensors/test_external_task_sensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py index 7e77fa60be527..97f5955ddada5 100644 --- a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py +++ b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py @@ -1185,8 +1185,6 @@ def func(dt, context): assert op._handle_execution_date_fn(context) == DEFAULT_DATE -@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Different test for AF 2") -@pytest.mark.usefixtures("testing_dag_bundle") def _api_server_error(status_code: int): """Build the error a task gets back when an execution API call fails.""" from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType @@ -1200,6 +1198,8 @@ def _api_server_error(status_code: int): ) +@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Different test for AF 2") +@pytest.mark.usefixtures("testing_dag_bundle") class TestExternalTaskSensorV3: def setup_method(self): # Create a mock for TaskInstance with get_ti_count method From bb6b6cbfaadc2883f1293fb03408d57ad32fba80 Mon Sep 17 00:00:00 2001 From: bingqin2 Date: Mon, 14 Sep 2026 09:03:30 -0500 Subject: [PATCH 3/3] Rely on the execution API's answer for missing tasks and task groups --- .../standard/sensors/external_task.py | 13 ++-- .../standard/triggers/external_task.py | 5 +- .../providers/standard/utils/sensor_helper.py | 68 ++++++++++++------- .../sensors/test_external_task_sensor.py | 58 +++++++++------- .../standard/triggers/test_external_task.py | 52 +++++++++----- 5 files changed, 121 insertions(+), 75 deletions(-) diff --git a/providers/standard/src/airflow/providers/standard/sensors/external_task.py b/providers/standard/src/airflow/providers/standard/sensors/external_task.py index a521495ffb2a0..70b221f71a22f 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py +++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py @@ -177,10 +177,10 @@ class ExternalTaskSensor(BaseSensorOperator): external_task_id is None), and immediately cease waiting if the external task or DAG does not exist (default value: False). On Airflow 3 a worker has no database access, so tasks and task groups are checked - through the execution API against each awaited Dag run once that run exists: a run's - task instances are created together with it from the run's own Dag version. Until the - run exists the sensor keeps waiting. Whether the Dag itself is registered is not checked - on Airflow 3. + through the execution API against each awaited Dag run once that run exists: the API + reports a task or task group that the run's Dag version does not define. Until the run + exists the sensor keeps waiting, and so does it with an API server that does not report + unknown tasks. Whether the Dag itself is registered is not checked on Airflow 3. :param poke_interval: polling period in seconds to check for the status :param poll_interval: (DEPRECATED) use ``poke_interval`` instead :param deferrable: Run sensor in deferrable mode @@ -569,9 +569,8 @@ def _check_for_existence_af3( Check that the awaited tasks or task group exist, through the execution API. A worker has no database access on Airflow 3, so unlike ``_check_for_existence`` this - relies on what the execution API exposes. A run's task instances are created together - with it from the run's own Dag version, so once an awaited run exists they answer - whether a task belongs to that run. While a run does not exist yet nothing can be + relies on what the execution API reports: a task or task group that the Dag version of + an existing awaited run does not define. While a run does not exist yet nothing can be concluded, so the check is repeated on later pokes until every awaited run has been seen. :param ti: the task instance running this sensor, used to reach the execution API diff --git a/providers/standard/src/airflow/providers/standard/triggers/external_task.py b/providers/standard/src/airflow/providers/standard/triggers/external_task.py index 0e64cdb652d0a..b86712c7b0425 100644 --- a/providers/standard/src/airflow/providers/standard/triggers/external_task.py +++ b/providers/standard/src/airflow/providers/standard/triggers/external_task.py @@ -52,8 +52,9 @@ class WorkflowTrigger(BaseTrigger): :param poke_interval: The interval (in seconds) for poking the external tasks. :param soft_fail: If True, the trigger will not fail the entire dag on external task failure. :param logical_dates: A list of logical dates for the external dag. - :param check_existence: If True, verify that the external tasks or task group exist in each - awaited Dag run once that run exists, and fire a ``not_found`` event otherwise. Airflow 3 only. + :param check_existence: If True, ask the execution API whether the external tasks or task group + exist in each awaited Dag run once that run exists, and fire a ``not_found`` event when it + reports them missing. Airflow 3 only. """ def __init__( diff --git a/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py b/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py index 6571064e5c43f..e554a6b05fa7f 100644 --- a/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py +++ b/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py @@ -158,6 +158,21 @@ def _get_count_by_matched_states( return count +def _not_found_message(error: Any) -> str | None: + """Return the execution API's message when ``error`` wraps a 404 response, ``None`` for any other error.""" + detail = error.error.detail or {} + if detail.get("status_code") != HTTPStatus.NOT_FOUND: + return None + # The supervisor forwards the server's JSON body under ``detail``; FastAPI nests an + # ``HTTPException`` detail under a ``detail`` key of its own. + payload = detail.get("detail") + while isinstance(payload, dict) and "message" not in payload and isinstance(payload.get("detail"), dict): + payload = payload["detail"] + if isinstance(payload, dict) and isinstance(payload.get("message"), str): + return payload["message"] + return str(detail.get("message") or "not found") + + def _check_external_task_existence( api: Any, *, @@ -170,14 +185,17 @@ def _check_external_task_existence( """ Verify that the awaited tasks or task group exist in the awaited Dag runs, through the execution API. - A Dag run's task instances are created together with the run, in the same transaction and - from the run's own Dag version, so once a run exists its task instances are the - version-accurate answer to whether a task belongs to it. Nothing can be concluded about a run - that does not exist yet, which is why the caller has to repeat the check until this function - returns True. - - :param api: an object exposing ``get_dr_count``, ``get_ti_count`` and ``get_task_states`` the - way ``RuntimeTaskInstance`` does: the running task instance, or the class itself. + The execution API answers ``task-instances/states`` with 404 when the Dag version an existing + run resolves to defines neither the requested tasks nor the requested task group and the run + has no task instance for them (apache/airflow#73086). A normal answer means they exist for + that run, even while their task instances have not been created yet, so nothing is inferred + from task-instance counts. Nothing can be concluded about a run that does not exist yet, + which is why the caller repeats the check until this function returns True. API servers + without that validation answer normally for unknown tasks, in which case the sensor keeps + waiting as it did before. + + :param api: an object exposing ``get_dr_count`` and ``get_task_states`` the way + ``RuntimeTaskInstance`` does: the running task instance, or the class itself. :param external_dag_id: The ID of the external Dag. :param external_task_ids: The task IDs that must exist in every awaited run. :param external_task_group_id: The task group ID that must exist in every awaited run. @@ -185,9 +203,8 @@ def _check_external_task_existence( :param run_ids: Run IDs identifying the awaited runs. :return: True once every awaited run exists and passed the check, False while at least one awaited run does not exist yet. - :raises ExternalTaskNotFoundError: when an existing run has no task instance for one of the tasks. - :raises ExternalTaskGroupNotFoundError: when the Dag has no such task group, or an existing run - has no task instance for any task of the group. + :raises ExternalTaskNotFoundError: when the API reports one of the tasks missing from an awaited run. + :raises ExternalTaskGroupNotFoundError: when the API reports the task group missing from an awaited run. """ from airflow.providers.standard.exceptions import ( ExternalTaskGroupNotFoundError, @@ -206,28 +223,29 @@ def _check_external_task_existence( if api.get_dr_count(dag_id=external_dag_id, **run_filter) == 0: all_runs_checked = False continue - for task_id in external_task_ids or (): - if api.get_ti_count(dag_id=external_dag_id, task_ids=[task_id], **run_filter) == 0: + + if external_task_ids: + try: + api.get_task_states(dag_id=external_dag_id, task_ids=list(external_task_ids), **run_filter) + except AirflowRuntimeError as e: + if (message := _not_found_message(e)) is None: + raise raise ExternalTaskNotFoundError( - f"The external task {task_id} in Dag {external_dag_id} does not exist for run {run_label}." - ) + f"The external tasks {list(external_task_ids)} in Dag {external_dag_id} " + f"do not all exist for run {run_label}: {message}" + ) from None if external_task_group_id: try: - run_id_task_state_map = api.get_task_states( + api.get_task_states( dag_id=external_dag_id, task_group_id=external_task_group_id, **run_filter ) except AirflowRuntimeError as e: - if (e.error.detail or {}).get("status_code") == HTTPStatus.NOT_FOUND: - raise ExternalTaskGroupNotFoundError( - f"The external task group '{external_task_group_id}' in Dag '{external_dag_id}' " - "does not exist." - ) from None - raise - if not any(run_id_task_state_map.values()): + if (message := _not_found_message(e)) is None: + raise raise ExternalTaskGroupNotFoundError( f"The external task group '{external_task_group_id}' in Dag '{external_dag_id}' " - f"does not exist for run {run_label}." - ) + f"does not exist for run {run_label}: {message}" + ) from None return all_runs_checked diff --git a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py index 97f5955ddada5..2bd064f4878ff 100644 --- a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py +++ b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py @@ -1185,7 +1185,7 @@ def func(dt, context): assert op._handle_execution_date_fn(context) == DEFAULT_DATE -def _api_server_error(status_code: int): +def _api_server_error(status_code: int, message: str | None = None): """Build the error a task gets back when an execution API call fails.""" from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType from airflow.sdk.execution_time.comms import ErrorResponse @@ -1193,7 +1193,12 @@ def _api_server_error(status_code: int): return AirflowRuntimeError( ErrorResponse( error=ErrorType.API_SERVER_ERROR, - detail={"status_code": status_code, "message": f"Server returned {status_code}"}, + detail={ + "status_code": status_code, + "message": f"Server returned {status_code}", + # the supervisor forwards the server's JSON body, which nests the HTTPException detail + "detail": {"detail": {"reason": "not_found", "message": message}} if message else None, + }, ) ) @@ -1665,24 +1670,22 @@ def test_check_existence_waits_for_run_then_checks_task_once(self, dag_maker): ) self.context["ti"].get_dr_count.side_effect = [0, 1] - # the existence probe (no ``states``) finds the task's instance; the state poll keeps waiting - self.context["ti"].get_ti_count.side_effect = lambda **kwargs: 0 if "states" in kwargs else 1 + # the API knows the task (no 404), even though no instance has reached an allowed state yet + self.context["ti"].get_task_states.return_value = {} + self.context["ti"].get_ti_count.return_value = 0 assert op.poke(self.context) is False # no run yet: nothing can be checked assert op.poke(self.context) is False # run exists: task verified, still waiting for its state assert op.poke(self.context) is False # already verified: no further run lookup assert self.context["ti"].get_dr_count.call_count == 2 - existence_probes = [ - c for c in self.context["ti"].get_ti_count.call_args_list if "states" not in c.kwargs - ] - assert existence_probes == [ - mock.call(dag_id="test_dag_parent", task_ids=["test_task"], logical_dates=[DEFAULT_DATE]), - ] + self.context["ti"].get_task_states.assert_called_once_with( + dag_id="test_dag_parent", task_ids=["test_task"], logical_dates=[DEFAULT_DATE] + ) @pytest.mark.execution_timeout(10) - def test_check_existence_task_not_found_in_run(self, dag_maker): - """A run that exists but has no task instance for the awaited task fails the sensor.""" + def test_check_existence_task_unknown_to_run_version(self, dag_maker): + """The execution API answers 404 for a task the awaited run's Dag version does not define.""" with dag_maker("test_dag_child"): op = ExternalTaskSensor( task_id="test_external_task_sensor_check", @@ -1692,12 +1695,14 @@ def test_check_existence_task_not_found_in_run(self, dag_maker): ) self.context["ti"].get_dr_count.return_value = 1 - self.context["ti"].get_ti_count.return_value = 0 + self.context["ti"].get_task_states.side_effect = _api_server_error( + 404, "Task missing_task not found in DAG test_dag_parent" + ) - with pytest.raises(ExternalTaskNotFoundError, match="missing_task"): + with pytest.raises(ExternalTaskNotFoundError, match="missing_task.*not found in DAG"): op.poke(self.context) - self.context["ti"].get_ti_count.assert_called_once_with( + self.context["ti"].get_task_states.assert_called_once_with( dag_id="test_dag_parent", task_ids=["missing_task"], logical_dates=[DEFAULT_DATE] ) @@ -1714,30 +1719,31 @@ def test_check_existence_task_not_found_ignores_soft_fail(self, dag_maker): ) self.context["ti"].get_dr_count.return_value = 1 - self.context["ti"].get_ti_count.return_value = 0 + self.context["ti"].get_task_states.side_effect = _api_server_error( + 404, "Task missing_task not found in DAG test_dag_parent" + ) with pytest.raises(ExternalTaskNotFoundError, match="missing_task"): op.execute(context=self.context) @pytest.mark.execution_timeout(10) - def test_check_existence_task_group_not_found_in_run(self, dag_maker): - """A run that exists without any task instance of the awaited group fails the sensor.""" + def test_check_existence_task_group_without_instances_keeps_waiting(self, dag_maker): + """A group the API knows but has no instances for yet is not reported missing.""" with dag_maker("test_dag_child"): op = ExternalTaskSensor( task_id="test_external_task_sensor_check", external_dag_id="test_dag_parent", - external_task_group_id="missing_group", + external_task_group_id="test_group", check_existence=True, ) self.context["ti"].get_dr_count.return_value = 1 self.context["ti"].get_task_states.return_value = {} - with pytest.raises(ExternalTaskGroupNotFoundError, match="missing_group"): - op.poke(self.context) + assert op.poke(self.context) is False - self.context["ti"].get_task_states.assert_called_once_with( - dag_id="test_dag_parent", task_group_id="missing_group", logical_dates=[DEFAULT_DATE] + self.context["ti"].get_task_states.assert_any_call( + dag_id="test_dag_parent", task_group_id="test_group", logical_dates=[DEFAULT_DATE] ) @pytest.mark.execution_timeout(10) @@ -1752,9 +1758,11 @@ def test_check_existence_task_group_unknown_to_dag(self, dag_maker): ) self.context["ti"].get_dr_count.return_value = 1 - self.context["ti"].get_task_states.side_effect = _api_server_error(404) + self.context["ti"].get_task_states.side_effect = _api_server_error( + 404, "Task group missing_group not found in DAG test_dag_parent" + ) - with pytest.raises(ExternalTaskGroupNotFoundError, match="missing_group"): + with pytest.raises(ExternalTaskGroupNotFoundError, match="missing_group.*not found in DAG"): op.poke(self.context) @pytest.mark.execution_timeout(10) diff --git a/providers/standard/tests/unit/standard/triggers/test_external_task.py b/providers/standard/tests/unit/standard/triggers/test_external_task.py index d8c0661d5ac8d..8d00b6a7bf6d3 100644 --- a/providers/standard/tests/unit/standard/triggers/test_external_task.py +++ b/providers/standard/tests/unit/standard/triggers/test_external_task.py @@ -39,6 +39,23 @@ key, value = next(iter(_DATES.items())) +def _api_server_error(status_code: int, message: str | None = None): + """Build the error a trigger gets back when an execution API call fails.""" + from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType + from airflow.sdk.execution_time.comms import ErrorResponse + + return AirflowRuntimeError( + ErrorResponse( + error=ErrorType.API_SERVER_ERROR, + detail={ + "status_code": status_code, + "message": f"Server returned {status_code}", + "detail": {"detail": {"reason": "not_found", "message": message}} if message else None, + }, + ) + ) + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 3") class TestWorkflowTrigger: DAG_ID = "external_task" @@ -421,12 +438,14 @@ def test_serialization_check_existence(self): assert kwargs["check_existence"] is True @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_dr_count") - @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_task_states") @pytest.mark.asyncio - async def test_task_workflow_trigger_task_not_found(self, mock_get_ti_count, mock_get_dr_count): - """A run that exists without the awaited task ends the wait with a not_found event.""" + async def test_task_workflow_trigger_task_not_found(self, mock_get_task_states, mock_get_dr_count): + """The API reporting the awaited task unknown to the run ends the wait with a not_found event.""" mock_get_dr_count.return_value = 1 - mock_get_ti_count.return_value = 0 # the run has no task instance for the awaited task + mock_get_task_states.side_effect = _api_server_error( + 404, f"Task {self.TASK_ID} not found in DAG {self.DAG_ID}" + ) trigger = WorkflowTrigger( external_dag_id=self.DAG_ID, logical_dates=[self.LOGICAL_DATE], @@ -441,7 +460,7 @@ async def test_task_workflow_trigger_task_not_found(self, mock_get_ti_count, moc assert result.payload["kind"] == "task" assert self.TASK_ID in result.payload["message"] mock_get_dr_count.assert_called_once_with(dag_id=self.DAG_ID, logical_dates=[self.LOGICAL_DATE]) - mock_get_ti_count.assert_any_call( + mock_get_task_states.assert_called_once_with( dag_id=self.DAG_ID, task_ids=[self.TASK_ID], logical_dates=[self.LOGICAL_DATE] ) with pytest.raises(StopAsyncIteration): @@ -454,9 +473,11 @@ async def test_task_workflow_trigger_task_not_found(self, mock_get_ti_count, moc async def test_task_workflow_trigger_task_group_not_found( self, mock_get_task_states, mock_get_ti_count, mock_get_dr_count ): - """A run that exists without any task of the awaited group ends the wait with a not_found event.""" + """The API reporting the group unknown to the run ends the wait with a not_found event.""" mock_get_dr_count.return_value = 1 - mock_get_task_states.return_value = {} + mock_get_task_states.side_effect = _api_server_error( + 404, f"Task group missing_group not found in DAG {self.DAG_ID}" + ) trigger = WorkflowTrigger( external_dag_id=self.DAG_ID, run_ids=[self.RUN_ID], @@ -477,22 +498,18 @@ async def test_task_workflow_trigger_task_group_not_found( await gen.__anext__() @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_dr_count") + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_task_states") @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") @mock.patch("asyncio.sleep") @pytest.mark.asyncio async def test_task_workflow_trigger_waits_for_run_before_checking( - self, mock_sleep, mock_get_ti_count, mock_get_dr_count + self, mock_sleep, mock_get_ti_count, mock_get_task_states, mock_get_dr_count ): """Nothing is concluded while the run is missing; once it exists the task is verified and the wait goes on.""" mock_get_dr_count.side_effect = [0, 1] - - def ti_count(**kwargs): - if "states" in kwargs: - # the state poll: the task reaches an allowed state only after the run exists - return 1 if mock_get_dr_count.call_count == 2 else 0 - return 1 # existence probes: the run has task instances, including the awaited task - - mock_get_ti_count.side_effect = ti_count + mock_get_task_states.return_value = {} # the API knows the task once the run exists + # the state poll: the task reaches an allowed state only after the run exists + mock_get_ti_count.side_effect = lambda **kwargs: 1 if mock_get_dr_count.call_count == 2 else 0 trigger = WorkflowTrigger( external_dag_id=self.DAG_ID, logical_dates=[self.LOGICAL_DATE], @@ -505,6 +522,9 @@ def ti_count(**kwargs): result = await gen.__anext__() assert result.payload == {"status": "success"} assert mock_get_dr_count.call_count == 2 + mock_get_task_states.assert_called_once_with( + dag_id=self.DAG_ID, task_ids=[self.TASK_ID], logical_dates=[self.LOGICAL_DATE] + ) mock_sleep.assert_awaited_once() with pytest.raises(StopAsyncIteration): await gen.__anext__()