From bf309a72978c7a3d61b89e78bb1203cc97b6f84f Mon Sep 17 00:00:00 2001 From: Cedrik Neumann Date: Wed, 30 Sep 2026 14:22:08 -0600 Subject: [PATCH] Add HttpExternalTaskSensor for Dags in remote Airflow deployments Waiting for a Dag, task group or task in another Airflow 3 deployment required copying most of ExternalTaskSensor and WorkflowTrigger, because their access to Airflow state was inlined into larger methods. Isolating that access in private methods lets the http provider reuse the existing sensor and trigger logic unchanged and only swap the data source for the remote deployment's REST API. Co-Authored-By: Claude --- providers/http/docs/index.rst | 30 ++- providers/http/docs/operators.rst | 36 +++ providers/http/provider.yaml | 2 + providers/http/pyproject.toml | 4 + .../providers/http/get_provider_info.py | 10 +- .../providers/http/sensors/external_task.py | 102 +++++++ .../providers/http/triggers/external_task.py | 223 +++++++++++++++ .../unit/http/sensors/test_external_task.py | 114 ++++++++ .../unit/http/triggers/test_external_task.py | 253 ++++++++++++++++++ .../standard/sensors/external_task.py | 86 +++--- .../standard/triggers/external_task.py | 54 ++-- .../sensors/test_external_task_sensor.py | 33 +++ .../standard/triggers/test_external_task.py | 32 +++ uv.lock | 8 +- 14 files changed, 932 insertions(+), 55 deletions(-) create mode 100644 providers/http/src/airflow/providers/http/sensors/external_task.py create mode 100644 providers/http/src/airflow/providers/http/triggers/external_task.py create mode 100644 providers/http/tests/unit/http/sensors/test_external_task.py create mode 100644 providers/http/tests/unit/http/triggers/test_external_task.py diff --git a/providers/http/docs/index.rst b/providers/http/docs/index.rst index 25455c9c365ed..f4bad86d85879 100644 --- a/providers/http/docs/index.rst +++ b/providers/http/docs/index.rst @@ -112,6 +112,25 @@ PIP package Version required ``pydantic`` ``>=2.11.0`` ========================================== ====================================== +Optional cross provider package dependencies +-------------------------------------------- + +Those are dependencies that might be needed in order to use all the features of the package. +You need to install the specified provider distributions in order to use them. + +You can install such cross-provider dependencies when installing from PyPI. For example: + +.. code-block:: bash + + pip install apache-airflow-providers-http[standard] + + +======================================================================================================== ============ +Dependent package Extra +======================================================================================================== ============ +`apache-airflow-providers-standard `_ ``standard`` +======================================================================================================== ============ + Optional dependencies --------------------- @@ -123,11 +142,12 @@ Install them when installing from PyPI. For example: pip install apache-airflow-providers-http[srv] -======= ==================== -Extra Dependencies -======= ==================== -``srv`` ``dnspython>=2.0.0`` -======= ==================== +============ ============================================= +Extra Dependencies +============ ============================================= +``srv`` ``dnspython>=2.0.0`` +``standard`` ``apache-airflow-providers-standard>=1.19.0`` +============ ============================================= Downloading official packages ----------------------------- diff --git a/providers/http/docs/operators.rst b/providers/http/docs/operators.rst index 1e75287742b1d..34559b2b6495a 100644 --- a/providers/http/docs/operators.rst +++ b/providers/http/docs/operators.rst @@ -44,6 +44,42 @@ This sensor can also be used in deferrable mode :start-after: [START howto_operator_http_http_sensor_check_deferrable] :end-before: [END howto_operator_http_http_sensor_check_deferrable] +.. _howto/operator:HttpExternalTaskSensor: + +HttpExternalTaskSensor +---------------------- + +Use the :class:`~airflow.providers.http.sensors.external_task.HttpExternalTaskSensor` to wait for a Dag, +task group, or task of a *remote* Airflow deployment. It accepts the same parameters and behaves the same +as :class:`~airflow.providers.standard.sensors.external_task.ExternalTaskSensor`, but queries the REST API +(v2) of the remote deployment instead of the local Airflow deployment. Both deployments need to run +Airflow 3, and waiting for a task group requires Airflow 3.2 or later on the remote deployment. +The sensor requires the ``standard`` extra (``apache-airflow-providers-http[standard]``). + +The ``http_conn_id`` connection points to the base URL of the remote Airflow API server, e.g. +``https://airflow.example.com``. The connection is used to authenticate as follows: + +* With login and password, the credentials are exchanged for a JWT access token via the ``/auth/token`` + endpoint of the remote deployment's auth manager. +* With a password but no login, the password is sent as a bearer token as-is, e.g. an API token issued by a + managed Airflow service. +* Without credentials, no authentication is added, but headers such as ``Authorization`` can be set via + the connection extra. + +The user on the remote deployment needs read access to the awaited Dag, its Dag runs and task instances. + +.. code-block:: python + + from airflow.providers.http.sensors.external_task import HttpExternalTaskSensor + + wait_for_remote_task = HttpExternalTaskSensor( + task_id="wait_for_remote_task", + http_conn_id="remote_airflow", + external_dag_id="remote_dag", + external_task_id="remote_task", + deferrable=True, + ) + .. _howto/operator:HttpOperator: HttpOperator diff --git a/providers/http/provider.yaml b/providers/http/provider.yaml index 85cc70bb2e816..7ebced03904ea 100644 --- a/providers/http/provider.yaml +++ b/providers/http/provider.yaml @@ -105,6 +105,7 @@ operators: sensors: - integration-name: Hypertext Transfer Protocol (HTTP) python-modules: + - airflow.providers.http.sensors.external_task - airflow.providers.http.sensors.http hooks: @@ -118,6 +119,7 @@ notifications: triggers: - integration-name: Hypertext Transfer Protocol (HTTP) python-modules: + - airflow.providers.http.triggers.external_task - airflow.providers.http.triggers.http connection-types: diff --git a/providers/http/pyproject.toml b/providers/http/pyproject.toml index 8447d9e8ffcd6..a7d9734e97586 100644 --- a/providers/http/pyproject.toml +++ b/providers/http/pyproject.toml @@ -77,6 +77,9 @@ dependencies = [ "srv" = [ "dnspython>=2.0.0", ] +"standard" = [ + "apache-airflow-providers-standard>=1.19.0", # use next version +] [dependency-groups] dev = [ @@ -84,6 +87,7 @@ dev = [ "apache-airflow-task-sdk", "apache-airflow-devel-common", "apache-airflow-providers-common-compat", + "apache-airflow-providers-standard", # Additional devel dependencies (do not remove this line and add extra development dependencies) ] diff --git a/providers/http/src/airflow/providers/http/get_provider_info.py b/providers/http/src/airflow/providers/http/get_provider_info.py index ee468a44fc6a2..21b5a8ac1a49e 100644 --- a/providers/http/src/airflow/providers/http/get_provider_info.py +++ b/providers/http/src/airflow/providers/http/get_provider_info.py @@ -44,7 +44,10 @@ def get_provider_info(): "sensors": [ { "integration-name": "Hypertext Transfer Protocol (HTTP)", - "python-modules": ["airflow.providers.http.sensors.http"], + "python-modules": [ + "airflow.providers.http.sensors.external_task", + "airflow.providers.http.sensors.http", + ], } ], "hooks": [ @@ -57,7 +60,10 @@ def get_provider_info(): "triggers": [ { "integration-name": "Hypertext Transfer Protocol (HTTP)", - "python-modules": ["airflow.providers.http.triggers.http"], + "python-modules": [ + "airflow.providers.http.triggers.external_task", + "airflow.providers.http.triggers.http", + ], } ], "connection-types": [ diff --git a/providers/http/src/airflow/providers/http/sensors/external_task.py b/providers/http/src/airflow/providers/http/sensors/external_task.py new file mode 100644 index 0000000000000..f72dcbb8fb0c2 --- /dev/null +++ b/providers/http/src/airflow/providers/http/sensors/external_task.py @@ -0,0 +1,102 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.http.hooks.http import HttpHook +from airflow.providers.http.triggers.external_task import HttpExternalTaskTrigger, _AirflowApiClient +from airflow.providers.standard.sensors.external_task import ExternalTaskSensor + +if TYPE_CHECKING: + import datetime + + from airflow.providers.common.compat.sdk import Context + + +class HttpExternalTaskSensor(ExternalTaskSensor): + """ + Waits for a Dag, task group, or task of a remote Airflow deployment to complete for a logical date. + + Behaves like :class:`~airflow.providers.standard.sensors.external_task.ExternalTaskSensor`, + but queries the REST API (v2) of the remote Airflow 3 deployment configured in the + :ref:`http connection` instead of the local Airflow deployment. + + Waiting for a task group requires the remote deployment to run Airflow 3.2 or later. + + .. seealso:: + For more information on how to use this sensor, take a look at the guide: + :ref:`howto/operator:HttpExternalTaskSensor` + + :param http_conn_id: The :ref:`http connection` of the remote + Airflow deployment. (templated) + + All other parameters are the same as those of + :class:`~airflow.providers.standard.sensors.external_task.ExternalTaskSensor`. + """ + + template_fields = [*ExternalTaskSensor.template_fields, "http_conn_id"] + # ExternalDagLink would point to the local deployment. + operator_extra_links = [] + + def __init__(self, *, http_conn_id: str = HttpHook.default_conn_name, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.http_conn_id = http_conn_id + + @cached_property + def _client(self) -> _AirflowApiClient: + return _AirflowApiClient(self.http_conn_id) + + def _get_dr_count( + self, context: Context, logical_dates: Sequence[datetime.datetime], states: list[str] + ) -> int: + return self._client.get_dr_count(self.external_dag_id, logical_dates, states) + + def _get_ti_count( + self, context: Context, logical_dates: Sequence[datetime.datetime], states: list[str] + ) -> int: + if TYPE_CHECKING: + assert self.external_task_ids + return self._client.get_ti_count(self.external_dag_id, self.external_task_ids, logical_dates, states) + + def _get_task_group_states( + self, context: Context, logical_dates: Sequence[datetime.datetime] + ) -> dict[str, dict[str, Any]]: + if TYPE_CHECKING: + assert self.external_task_group_id + return self._client.get_task_group_states( + self.external_dag_id, self.external_task_group_id, logical_dates + ) + + def _get_trigger( + self, context: Context, logical_dates: Sequence[datetime.datetime] + ) -> HttpExternalTaskTrigger: + return HttpExternalTaskTrigger( + http_conn_id=self.http_conn_id, + external_dag_id=self.external_dag_id, + external_task_group_id=self.external_task_group_id, + external_task_ids=self.external_task_ids, + allowed_states=self.allowed_states, + failed_states=self.failed_states, + skipped_states=self.skipped_states, + poke_interval=self.poke_interval, + soft_fail=self.soft_fail, + logical_dates=list(logical_dates), + ) diff --git a/providers/http/src/airflow/providers/http/triggers/external_task.py b/providers/http/src/airflow/providers/http/triggers/external_task.py new file mode 100644 index 0000000000000..6973100cbb29b --- /dev/null +++ b/providers/http/src/airflow/providers/http/triggers/external_task.py @@ -0,0 +1,223 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from collections.abc import Collection, Iterable, Iterator +from functools import cached_property +from typing import TYPE_CHECKING, Any +from urllib.parse import quote + +from asgiref.sync import sync_to_async +from packaging.version import Version + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +from airflow.providers.http.hooks.http import HttpHook +from airflow.providers.http.version_compat import AIRFLOW_V_3_0_PLUS + +if not AIRFLOW_V_3_0_PLUS: + raise AirflowOptionalProviderFeatureException("Waiting for a remote Airflow deployment needs Airflow 3+.") + +from airflow.providers.standard.triggers.external_task import WorkflowTrigger + +if TYPE_CHECKING: + from datetime import datetime + + from requests import Response, Session + +# The ``task_group_id`` filter of the task instances endpoint was added in Airflow 3.2.0. +# Older API servers silently ignore it and would return the task instances of all tasks. +MIN_TASK_GROUP_FILTER_VERSION = Version("3.2.0") + + +class _AirflowApiClient: + """ + Query the REST API (v2) of the remote Airflow deployment of an HTTP connection. + + Login and password of the connection are exchanged for a JWT access token of the remote auth + manager, a password without login is used as bearer token as-is. + """ + + page_limit = 100 + + def __init__(self, http_conn_id: str) -> None: + self.http_conn_id = http_conn_id + self._hook = HttpHook(method="GET", http_conn_id=http_conn_id) + self._session: Session | None = None + self._uses_access_token = False + self._remote_version: Version | None = None + + def _create_session(self) -> Session: + session = self._hook.get_conn() + # Airflow 3 API servers reject the basic auth HttpHook derives from the connection login. + session.auth = None + connection = self._hook.get_connection(self.http_conn_id) + self._uses_access_token = bool(connection.login) + if connection.login: + response = session.post( + self._hook.url_from_endpoint("auth/token"), + json={"username": connection.login, "password": connection.password}, + timeout=self._hook.merged_extra.get("timeout"), + ) + response.raise_for_status() + session.headers["Authorization"] = f"Bearer {response.json()['access_token']}" + elif connection.password: + session.headers["Authorization"] = f"Bearer {connection.password}" + return session + + def _request(self, method: str, endpoint: str, **kwargs: Any) -> Any: + if self._session is None: + self._session = self._create_session() + url = self._hook.url_from_endpoint(f"api/v2/{endpoint}") + timeout = self._hook.merged_extra.get("timeout") + response: Response = self._session.request(method, url, timeout=timeout, **kwargs) + if response.status_code == 401 and self._uses_access_token: + # Access tokens expire, which long-running sensors and triggers outlive. + self._session = self._create_session() + response = self._session.request(method, url, timeout=timeout, **kwargs) + response.raise_for_status() + return response.json() + + def get_dr_count( + self, dag_id: str, logical_dates: Iterable[datetime], states: Collection[str] | None + ) -> int: + return sum( + self._request( + "GET", + f"dags/{quote(dag_id, safe='')}/dagRuns", + params={ + "logical_date_gte": logical_date.isoformat(), + "logical_date_lte": logical_date.isoformat(), + "state": list(states or []), + "limit": 1, + }, + )["total_entries"] + for logical_date in logical_dates + ) + + def get_ti_count( + self, + dag_id: str, + task_ids: Collection[str], + logical_dates: Iterable[datetime], + states: Collection[str] | None, + ) -> int: + return sum( + self._request( + "POST", + "dags/~/dagRuns/~/taskInstances/list", + json={ + "dag_ids": [dag_id], + "task_ids": list(task_ids), + "state": list(states) if states else None, + "logical_date_gte": logical_date.isoformat(), + "logical_date_lte": logical_date.isoformat(), + "page_limit": 1, + }, + )["total_entries"] + for logical_date in logical_dates + ) + + def get_task_group_states( + self, dag_id: str, task_group_id: str, logical_dates: Iterable[datetime] + ) -> dict[str, dict[str, Any]]: + version = self._get_remote_version() + if version.release < MIN_TASK_GROUP_FILTER_VERSION.release: + raise ValueError( + f"Waiting for a task group requires the remote Airflow deployment to run Airflow " + f"{MIN_TASK_GROUP_FILTER_VERSION} or later, but it runs Airflow {version}." + ) + task_states: dict[str, dict[str, Any]] = {} + for logical_date in logical_dates: + for ti in self._iter_task_group_instances(dag_id, task_group_id, logical_date): + # Keyed like the Execution API's task states to reuse the same state matching. + key = ti["task_id"] if ti["map_index"] < 0 else f"{ti['task_id']}_{ti['map_index']}" + task_states.setdefault(ti["dag_run_id"], {})[key] = ti["state"] + return task_states + + def _iter_task_group_instances( + self, dag_id: str, task_group_id: str, logical_date: datetime + ) -> Iterator[dict[str, Any]]: + offset = 0 + while True: + page = self._request( + "GET", + f"dags/{quote(dag_id, safe='')}/dagRuns/~/taskInstances", + params={ + "task_group_id": task_group_id, + "logical_date_gte": logical_date.isoformat(), + "logical_date_lte": logical_date.isoformat(), + "order_by": "id", + "limit": self.page_limit, + "offset": offset, + }, + ) + yield from page["task_instances"] + offset += len(page["task_instances"]) + if not page["task_instances"] or offset >= page["total_entries"]: + return + + def _get_remote_version(self) -> Version: + if self._remote_version is None: + self._remote_version = Version(self._request("GET", "version")["version"]) + return self._remote_version + + +class HttpExternalTaskTrigger(WorkflowTrigger): + """ + Wait for a Dag, task group or task of a remote Airflow deployment to reach a state. + + Behaves like :class:`~airflow.providers.standard.triggers.external_task.WorkflowTrigger`, + but queries the REST API (v2) of the remote Airflow 3 deployment configured in the HTTP connection. + Only ``logical_dates`` are supported to select the remote Dag runs. + + :param http_conn_id: :ref:`http connection` of the remote Airflow deployment. + + All other parameters are the same as those of + :class:`~airflow.providers.standard.triggers.external_task.WorkflowTrigger`. + """ + + def __init__(self, *, http_conn_id: str, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.http_conn_id = http_conn_id + + def serialize(self) -> tuple[str, dict[str, Any]]: + classpath, data = super().serialize() + return classpath, {**data, "http_conn_id": self.http_conn_id} + + @cached_property + def _client(self) -> _AirflowApiClient: + return _AirflowApiClient(self.http_conn_id) + + async def _get_dr_count(self, states: Collection[str] | None) -> int: + return await sync_to_async(self._client.get_dr_count)( + self.external_dag_id, self.logical_dates or [], states + ) + + async def _get_ti_count(self, states: Collection[str] | None) -> int: + if TYPE_CHECKING: + assert self.external_task_ids + return await sync_to_async(self._client.get_ti_count)( + self.external_dag_id, self.external_task_ids, self.logical_dates or [], states + ) + + async def _get_task_group_states(self) -> dict[str, dict[str, Any]]: + if TYPE_CHECKING: + assert self.external_task_group_id + return await sync_to_async(self._client.get_task_group_states)( + self.external_dag_id, self.external_task_group_id, self.logical_dates or [] + ) diff --git a/providers/http/tests/unit/http/sensors/test_external_task.py b/providers/http/tests/unit/http/sensors/test_external_task.py new file mode 100644 index 0000000000000..955b541281c5b --- /dev/null +++ b/providers/http/tests/unit/http/sensors/test_external_task.py @@ -0,0 +1,114 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from unittest import mock + +import pytest + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +if not AIRFLOW_V_3_0_PLUS: + pytest.skip("Waiting for a remote Airflow deployment needs Airflow 3+", allow_module_level=True) + +from airflow.providers.common.compat.sdk import TaskDeferred +from airflow.providers.http.sensors.external_task import HttpExternalTaskSensor +from airflow.providers.http.triggers.external_task import HttpExternalTaskTrigger +from airflow.providers.standard.exceptions import ExternalTaskFailedError +from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + +LOGICAL_DATE = datetime(2026, 1, 2, tzinfo=timezone.utc) + + +@pytest.fixture +def client(): + with mock.patch( + "airflow.providers.http.sensors.external_task._AirflowApiClient", autospec=True + ) as client_class: + yield client_class.return_value + + +@pytest.fixture +def context(): + return {"logical_date": LOGICAL_DATE, "ti": mock.create_autospec(RuntimeTaskInstance, instance=True)} + + +def create_sensor(**kwargs) -> HttpExternalTaskSensor: + return HttpExternalTaskSensor( + task_id="wait", http_conn_id="remote_airflow", external_dag_id="remote_dag", **kwargs + ) + + +class TestHttpExternalTaskSensor: + def test_attributes(self): + sensor = create_sensor() + + assert sensor._client.http_conn_id == "remote_airflow" + assert "http_conn_id" in sensor.template_fields + assert not sensor.operator_extra_links + + @pytest.mark.parametrize(("count", "expected"), [(1, True), (0, False)]) + def test_poke_dag(self, client, context, count, expected): + client.get_dr_count.return_value = count + + assert create_sensor(execution_delta=timedelta(days=1)).poke(context) is expected + client.get_dr_count.assert_called_once_with( + "remote_dag", [LOGICAL_DATE - timedelta(days=1)], ["success"] + ) + context["ti"].get_dr_count.assert_not_called() + + def test_poke_tasks(self, client, context): + client.get_ti_count.return_value = 1 + + with pytest.raises(ExternalTaskFailedError): + create_sensor(external_task_ids=["t1", "t2"], failed_states=["failed"]).poke(context) + client.get_ti_count.assert_called_once_with("remote_dag", ["t1", "t2"], [LOGICAL_DATE], ["failed"]) + context["ti"].get_ti_count.assert_not_called() + + @pytest.mark.parametrize( + ("task_states", "expected"), + [ + ({"run_1": {"g.a": "success", "g.b_0": "success"}}, True), + ({"run_1": {"g.a": "success", "g.b_0": "running"}}, False), + ], + ids=["all_allowed", "partially_allowed"], + ) + def test_poke_task_group(self, client, context, task_states, expected): + client.get_task_group_states.return_value = task_states + + assert create_sensor(external_task_group_id="g").poke(context) is expected + client.get_task_group_states.assert_called_once_with("remote_dag", "g", [LOGICAL_DATE]) + context["ti"].get_task_states.assert_not_called() + + def test_execute_deferrable(self, context): + sensor = create_sensor( + external_task_ids=["t1"], failed_states=["failed"], deferrable=True, poke_interval=30 + ) + + with pytest.raises(TaskDeferred) as deferred: + sensor.execute(context) + + trigger = deferred.value.trigger + assert isinstance(trigger, HttpExternalTaskTrigger) + assert trigger.http_conn_id == "remote_airflow" + assert trigger.external_dag_id == "remote_dag" + assert trigger.external_task_ids == ["t1"] + assert trigger.failed_states == ["failed"] + assert trigger.logical_dates == [LOGICAL_DATE] + assert trigger.poke_interval == 30 diff --git a/providers/http/tests/unit/http/triggers/test_external_task.py b/providers/http/tests/unit/http/triggers/test_external_task.py new file mode 100644 index 0000000000000..2faa40231674a --- /dev/null +++ b/providers/http/tests/unit/http/triggers/test_external_task.py @@ -0,0 +1,253 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from datetime import datetime, timezone +from urllib.parse import parse_qs, urlsplit + +import pytest +import requests + +from airflow.models import Connection + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +if not AIRFLOW_V_3_0_PLUS: + pytest.skip("Waiting for a remote Airflow deployment needs Airflow 3+", allow_module_level=True) + +from airflow.providers.http.triggers.external_task import HttpExternalTaskTrigger, _AirflowApiClient +from airflow.triggers.base import TriggerEvent + +BASE_URL = "https://airflow.example.com" +API_URL = f"{BASE_URL}/api/v2" +DATE_1 = datetime(2026, 1, 1, tzinfo=timezone.utc) +DATE_2 = datetime(2026, 1, 2, tzinfo=timezone.utc) + + +def get_query(request) -> dict[str, list[str]]: + # requests_mock's ``qs`` lowercases values, which mangles ISO timestamps. + return parse_qs(urlsplit(request.url).query) + + +@pytest.fixture(autouse=True) +def setup_connections(create_connection_without_db): + create_connection_without_db( + Connection(conn_id="with_login", conn_type="http", host=BASE_URL, login="user", password="secret") + ) + create_connection_without_db( + Connection(conn_id="with_token", conn_type="http", host=BASE_URL, password="api-token") + ) + create_connection_without_db(Connection(conn_id="anonymous", conn_type="http", host=BASE_URL)) + + +@pytest.fixture +def token_mock(requests_mock): + return requests_mock.post( + f"{BASE_URL}/auth/token", + [ + {"json": {"access_token": "jwt-1"}, "status_code": 201}, + {"json": {"access_token": "jwt-2"}, "status_code": 201}, + ], + ) + + +@pytest.fixture +def client(token_mock): + return _AirflowApiClient("with_login") + + +class TestAirflowApiClientAuth: + def test_login_is_exchanged_for_access_token_once(self, requests_mock, token_mock): + requests_mock.get(f"{API_URL}/version", json={"version": "3.2.0"}) + client = _AirflowApiClient("with_login") + + client._request("GET", "version") + client._request("GET", "version") + + assert token_mock.call_count == 1 + assert token_mock.last_request.json() == {"username": "user", "password": "secret"} + assert "Authorization" not in token_mock.last_request.headers + assert requests_mock.last_request.headers["Authorization"] == "Bearer jwt-1" + + def test_expired_access_token_is_refreshed(self, requests_mock, token_mock): + version_mock = requests_mock.get( + f"{API_URL}/version", [{"status_code": 401}, {"json": {"version": "3.2.0"}}] + ) + + assert _AirflowApiClient("with_login")._request("GET", "version") == {"version": "3.2.0"} + assert token_mock.call_count == 2 + assert version_mock.last_request.headers["Authorization"] == "Bearer jwt-2" + + @pytest.mark.parametrize( + ("conn_id", "expected_header"), [("with_token", "Bearer api-token"), ("anonymous", None)] + ) + def test_without_login(self, requests_mock, conn_id, expected_header): + version_mock = requests_mock.get(f"{API_URL}/version", [{"status_code": 401}]) + + with pytest.raises(requests.HTTPError, match="401"): + _AirflowApiClient(conn_id)._request("GET", "version") + + assert version_mock.call_count == 1 + assert version_mock.last_request.headers.get("Authorization") == expected_header + + +class TestAirflowApiClient: + def test_get_dr_count(self, requests_mock, client): + dr_mock = requests_mock.get( + f"{API_URL}/dags/my_dag/dagRuns", + [{"json": {"dag_runs": [], "total_entries": 1}}, {"json": {"dag_runs": [], "total_entries": 0}}], + ) + + assert client.get_dr_count("my_dag", [DATE_1, DATE_2], ["success", "failed"]) == 1 + assert [get_query(r) for r in dr_mock.request_history] == [ + { + "logical_date_gte": [date.isoformat()], + "logical_date_lte": [date.isoformat()], + "state": ["success", "failed"], + "limit": ["1"], + } + for date in (DATE_1, DATE_2) + ] + + def test_get_ti_count(self, requests_mock, client): + ti_mock = requests_mock.post( + f"{API_URL}/dags/~/dagRuns/~/taskInstances/list", + [ + {"json": {"task_instances": [], "total_entries": 3}}, + {"json": {"task_instances": [], "total_entries": 2}}, + ], + ) + + assert client.get_ti_count("my_dag", ["t1", "t2"], [DATE_1, DATE_2], ["success"]) == 5 + assert [r.json() for r in ti_mock.request_history] == [ + { + "dag_ids": ["my_dag"], + "task_ids": ["t1", "t2"], + "state": ["success"], + "logical_date_gte": date.isoformat(), + "logical_date_lte": date.isoformat(), + "page_limit": 1, + } + for date in (DATE_1, DATE_2) + ] + + def test_get_task_group_states(self, requests_mock, client, monkeypatch): + monkeypatch.setattr(_AirflowApiClient, "page_limit", 2) + requests_mock.get(f"{API_URL}/version", json={"version": "3.2.0"}) + + def ti(task_id, map_index, state): + return {"dag_run_id": "run_1", "task_id": task_id, "map_index": map_index, "state": state} + + ti_mock = requests_mock.get( + f"{API_URL}/dags/my_dag/dagRuns/~/taskInstances", + [ + { + "json": { + "task_instances": [ti("g.a", -1, "success"), ti("g.b", 0, "failed")], + "total_entries": 3, + } + }, + {"json": {"task_instances": [ti("g.b", 1, None)], "total_entries": 3}}, + ], + ) + + assert client.get_task_group_states("my_dag", "g", [DATE_1]) == { + "run_1": {"g.a": "success", "g.b_0": "failed", "g.b_1": None} + } + assert [get_query(r)["offset"] for r in ti_mock.request_history] == [["0"], ["2"]] + assert get_query(ti_mock.last_request) == { + "task_group_id": ["g"], + "logical_date_gte": [DATE_1.isoformat()], + "logical_date_lte": [DATE_1.isoformat()], + "order_by": ["id"], + "limit": ["2"], + "offset": ["2"], + } + + def test_get_task_group_states_requires_airflow_3_2(self, requests_mock, client): + requests_mock.get(f"{API_URL}/version", json={"version": "3.1.8"}) + + with pytest.raises( + ValueError, match=r"requires the remote Airflow deployment to run Airflow 3\.2\.0" + ): + client.get_task_group_states("my_dag", "g", [DATE_1]) + + +def create_trigger(**kwargs) -> HttpExternalTaskTrigger: + return HttpExternalTaskTrigger( + http_conn_id="anonymous", + external_dag_id="my_dag", + logical_dates=[DATE_1], + allowed_states=["success"], + poke_interval=5, + **kwargs, + ) + + +async def get_first_event(trigger: HttpExternalTaskTrigger) -> TriggerEvent: + return await anext(aiter(trigger.run())) + + +class TestHttpExternalTaskTrigger: + def test_serialize(self): + classpath, kwargs = create_trigger(external_task_ids=["t1"]).serialize() + + assert classpath == "airflow.providers.http.triggers.external_task.HttpExternalTaskTrigger" + assert kwargs["http_conn_id"] == "anonymous" + assert HttpExternalTaskTrigger(**kwargs).serialize() == (classpath, kwargs) + + @pytest.mark.asyncio + async def test_run_dag(self, requests_mock): + dr_mock = requests_mock.get( + f"{API_URL}/dags/my_dag/dagRuns", json={"dag_runs": [], "total_entries": 1} + ) + + assert await get_first_event(create_trigger()) == TriggerEvent({"status": "success"}) + assert get_query(dr_mock.last_request)["state"] == ["success"] + + @pytest.mark.asyncio + async def test_run_tasks(self, requests_mock): + ti_mock = requests_mock.post( + f"{API_URL}/dags/~/dagRuns/~/taskInstances/list", json={"task_instances": [], "total_entries": 2} + ) + + event = await get_first_event( + create_trigger(external_task_ids=["t1", "t2"], failed_states=["failed"]) + ) + + assert event == TriggerEvent({"status": "failed"}) + assert ti_mock.last_request.json()["task_ids"] == ["t1", "t2"] + assert ti_mock.last_request.json()["state"] == ["failed"] + + @pytest.mark.asyncio + async def test_run_task_group(self, requests_mock): + requests_mock.get(f"{API_URL}/version", json={"version": "3.2.0"}) + ti_mock = requests_mock.get( + f"{API_URL}/dags/my_dag/dagRuns/~/taskInstances", + json={ + "task_instances": [ + {"dag_run_id": "run_1", "task_id": "g.a", "map_index": -1, "state": "success"} + ], + "total_entries": 1, + }, + ) + + assert await get_first_event(create_trigger(external_task_group_id="g")) == TriggerEvent( + {"status": "success"} + ) + assert get_query(ti_mock.last_request)["task_group_id"] == ["g"] 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..b7a255545cd8c 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py +++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py @@ -64,6 +64,7 @@ from sqlalchemy.orm import Session from airflow.providers.common.compat.sdk import Context, TaskInstanceKey + from airflow.triggers.base import BaseTrigger class ExternalDagLink(BaseOperatorLink): @@ -352,28 +353,14 @@ def _poke_af3(self, context: Context, dttm_filter: Sequence[datetime.datetime]) from airflow.providers.standard.utils.sensor_helper import _get_count_by_matched_states self._has_checked_existence = True - ti = context["ti"] def _get_count(states: list[str]) -> int: if self.external_task_ids: - return ti.get_ti_count( - dag_id=self.external_dag_id, - task_ids=list(self.external_task_ids), - logical_dates=list(dttm_filter), - states=states, - ) + return self._get_ti_count(context, dttm_filter, states) if self.external_task_group_id: - run_id_task_state_map = ti.get_task_states( - dag_id=self.external_dag_id, - task_group_id=self.external_task_group_id, - logical_dates=list(dttm_filter), - ) + run_id_task_state_map = self._get_task_group_states(context, dttm_filter) return _get_count_by_matched_states(run_id_task_state_map, states) - return ti.get_dr_count( - dag_id=self.external_dag_id, - logical_dates=list(dttm_filter), - states=states, - ) + return self._get_dr_count(context, dttm_filter, states) if self.failed_states: count = _get_count(self.failed_states) @@ -389,6 +376,41 @@ def _get_count(states: list[str]) -> int: count_allowed = self._calculate_count(count, dttm_filter) return count_allowed == len(dttm_filter) + def _get_dr_count( + self, context: Context, logical_dates: Sequence[datetime.datetime], states: list[str] + ) -> int: + """Count the external Dag runs with any of the logical dates and states.""" + return context["ti"].get_dr_count( + dag_id=self.external_dag_id, + logical_dates=list(logical_dates), + states=states, + ) + + def _get_ti_count( + self, context: Context, logical_dates: Sequence[datetime.datetime], states: list[str] + ) -> int: + """Count the external task instances with any of the logical dates and states.""" + if TYPE_CHECKING: + assert self.external_task_ids + return context["ti"].get_ti_count( + dag_id=self.external_dag_id, + task_ids=list(self.external_task_ids), + logical_dates=list(logical_dates), + states=states, + ) + + def _get_task_group_states( + self, context: Context, logical_dates: Sequence[datetime.datetime] + ) -> dict[str, dict[str, typing.Any]]: + """Get the states of the external task group's task instances by run ID and task key.""" + if TYPE_CHECKING: + assert self.external_task_group_id + return context["ti"].get_task_states( + dag_id=self.external_dag_id, + task_group_id=self.external_task_group_id, + logical_dates=list(logical_dates), + ) + def _calculate_count(self, count: int, dttm_filter: Sequence[datetime.datetime]) -> float | int: """Calculate the normalized count based on the type of check.""" if self.external_task_ids: @@ -476,19 +498,7 @@ def execute(self, context: Context) -> None: if AIRFLOW_V_3_0_PLUS: self.defer( timeout=datetime.timedelta(seconds=timeout_value) if timeout_value else None, - trigger=WorkflowTrigger( - external_dag_id=self.external_dag_id, - external_task_group_id=self.external_task_group_id, - external_task_ids=self.external_task_ids, - allowed_states=self.allowed_states, - failed_states=self.failed_states, - skipped_states=self.skipped_states, - poke_interval=self.poke_interval, - soft_fail=self.soft_fail, - logical_dates=list(dttm_filter), - run_ids=None, - execution_dates=None, - ), + trigger=self._get_trigger(context, dttm_filter), method_name="execute_complete", ) else: @@ -517,6 +527,22 @@ def execute(self, context: Context) -> None: method_name="execute_complete", ) + def _get_trigger(self, context: Context, logical_dates: Sequence[datetime.datetime]) -> BaseTrigger: + """Create the trigger to defer to in deferrable mode.""" + return WorkflowTrigger( + external_dag_id=self.external_dag_id, + external_task_group_id=self.external_task_group_id, + external_task_ids=self.external_task_ids, + allowed_states=self.allowed_states, + failed_states=self.failed_states, + skipped_states=self.skipped_states, + poke_interval=self.poke_interval, + soft_fail=self.soft_fail, + logical_dates=list(logical_dates), + run_ids=None, + execution_dates=None, + ) + def execute_complete(self, context: Context, event: dict[str, typing.Any] | None = None) -> None: """Execute when the trigger fires - return immediately.""" if event is None: 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..949fee0782465 100644 --- a/providers/standard/src/airflow/providers/standard/triggers/external_task.py +++ b/providers/standard/src/airflow/providers/standard/triggers/external_task.py @@ -99,7 +99,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: else: data["execution_dates"] = self.execution_dates - return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data + return f"{type(self).__module__}.{type(self).__qualname__}", data async def run(self) -> typing.AsyncIterator[TriggerEvent]: """Check periodically tasks, task group or dag status.""" @@ -132,36 +132,56 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]: 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 if self.external_task_ids: - count = await sync_to_async(RuntimeTaskInstance.get_ti_count)( - dag_id=self.external_dag_id, - task_ids=list(self.external_task_ids), - logical_dates=self.logical_dates, - run_ids=self.run_ids, - states=list(states) if states else None, - ) + count = await self._get_ti_count(states) return int(count / len(self.external_task_ids)) if self.external_task_group_id: - run_id_task_state_map = await sync_to_async(RuntimeTaskInstance.get_task_states)( - dag_id=self.external_dag_id, - task_group_id=self.external_task_group_id, - logical_dates=self.logical_dates, - run_ids=self.run_ids, - ) + run_id_task_state_map = await self._get_task_group_states() count = await sync_to_async(_get_count_by_matched_states)( run_id_task_state_map=run_id_task_state_map, states=states or [], ) return count - count = await sync_to_async(RuntimeTaskInstance.get_dr_count)( + return await self._get_dr_count(states) + + async def _get_dr_count(self, states: Collection[str] | None) -> int: + """Count the external Dag runs with any of the run IDs or logical dates and states.""" + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + return await sync_to_async(RuntimeTaskInstance.get_dr_count)( dag_id=self.external_dag_id, logical_dates=self.logical_dates, run_ids=self.run_ids, states=list(states) if states else None, ) - return count + + async def _get_ti_count(self, states: Collection[str] | None) -> int: + """Count the external task instances with any of the run IDs or logical dates and states.""" + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + if typing.TYPE_CHECKING: + assert self.external_task_ids + return await sync_to_async(RuntimeTaskInstance.get_ti_count)( + dag_id=self.external_dag_id, + task_ids=list(self.external_task_ids), + logical_dates=self.logical_dates, + run_ids=self.run_ids, + states=list(states) if states else None, + ) + + async def _get_task_group_states(self) -> dict[str, dict[str, typing.Any]]: + """Get the states of the external task group's task instances by run ID and task key.""" + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + if typing.TYPE_CHECKING: + assert self.external_task_group_id + return await sync_to_async(RuntimeTaskInstance.get_task_states)( + dag_id=self.external_dag_id, + task_group_id=self.external_task_group_id, + logical_dates=self.logical_dates, + run_ids=self.run_ids, + ) @sync_to_async def _get_count(self, states: Collection[str] | None) -> int: 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..5d83bb61b7d23 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 @@ -1560,6 +1560,39 @@ def test_deferrable_poke_interval_passed_to_trigger(self, dag_maker): assert isinstance(exc.value.trigger, WorkflowTrigger) assert exc.value.trigger.poke_interval == 30 + @pytest.mark.parametrize( + ("kwargs", "method", "return_value", "extra_args"), + [ + ({}, "_get_dr_count", 1, [["success"]]), + ({"external_task_ids": ["t1", "t2"]}, "_get_ti_count", 2, [["success"]]), + ({"external_task_group_id": "g"}, "_get_task_group_states", {"run_id": {"g.t1": "success"}}, []), + ], + ids=["dag", "tasks", "task_group"], + ) + def test_poke_uses_overridable_state_access(self, kwargs, method, return_value, extra_args): + op = ExternalTaskSensor(task_id="sensor", external_dag_id="test_dag_parent", **kwargs) + + with mock.patch.object(ExternalTaskSensor, method, autospec=True, return_value=return_value) as m: + assert op.poke(context=self.context) + + m.assert_called_once_with(op, self.context, [DEFAULT_DATE], *extra_args) + self.context["ti"].get_dr_count.assert_not_called() + self.context["ti"].get_ti_count.assert_not_called() + self.context["ti"].get_task_states.assert_not_called() + + def test_execute_defers_to_overridable_trigger(self): + op = ExternalTaskSensor(task_id="sensor", external_dag_id="test_dag_parent", deferrable=True) + trigger = mock.MagicMock(spec=WorkflowTrigger) + + with ( + mock.patch.object(ExternalTaskSensor, "_get_trigger", autospec=True, return_value=trigger) as m, + pytest.raises(TaskDeferred) as exc, + ): + op.execute(context=self.context) + + m.assert_called_once_with(op, self.context, [DEFAULT_DATE]) + assert exc.value.trigger is trigger + @pytest.mark.execution_timeout(10) def test_external_task_sensor_only_dag_id(self, dag_maker): """Test that the sensor works correctly when only external_dag_id is provided.""" 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..ebb495805cbf4 100644 --- a/providers/standard/tests/unit/standard/triggers/test_external_task.py +++ b/providers/standard/tests/unit/standard/triggers/test_external_task.py @@ -409,6 +409,38 @@ def test_serialization(self): "soft_fail": False, } + def test_serialization_of_subclass(self): + classpath, _ = _WorkflowTriggerSubclass(external_dag_id=self.DAG_ID).serialize() + + assert classpath == f"{__name__}._WorkflowTriggerSubclass" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("kwargs", "method", "return_value"), + [ + ({}, "_get_dr_count", 1), + ({"external_task_ids": ["t1", "t2"]}, "_get_ti_count", 2), + ({"external_task_group_id": "g"}, "_get_task_group_states", {"run_id": {"g.t1": "success"}}), + ], + ids=["dag", "tasks", "task_group"], + ) + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance", autospec=True) + async def test_run_uses_overridable_state_access(self, mock_ti, kwargs, method, return_value): + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, run_ids=[self.RUN_ID], allowed_states=["success"], **kwargs + ) + + with mock.patch.object(WorkflowTrigger, method, autospec=True, return_value=return_value) as m: + event = await anext(aiter(trigger.run())) + + assert event == TriggerEvent({"status": "success"}) + m.assert_awaited_once() + assert not mock_ti.method_calls + + +class _WorkflowTriggerSubclass(WorkflowTrigger): + pass + @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 2") class TestWorkflowTriggerAF2: diff --git a/uv.lock b/uv.lock index ecf33af68169a..1d6b0bdd9d9db 100644 --- a/uv.lock +++ b/uv.lock @@ -6240,12 +6240,16 @@ dependencies = [ srv = [ { name = "dnspython" }, ] +standard = [ + { name = "apache-airflow-providers-standard" }, +] [package.dev-dependencies] dev = [ { name = "apache-airflow" }, { name = "apache-airflow-devel-common" }, { name = "apache-airflow-providers-common-compat" }, + { name = "apache-airflow-providers-standard" }, { name = "apache-airflow-task-sdk" }, ] docs = [ @@ -6257,6 +6261,7 @@ requires-dist = [ { name = "aiohttp", specifier = ">=3.14.0" }, { name = "apache-airflow", editable = "." }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, + { name = "apache-airflow-providers-standard", marker = "extra == 'standard'", editable = "providers/standard" }, { name = "asgiref", marker = "python_full_version < '3.14'", specifier = ">=2.3.0" }, { name = "asgiref", marker = "python_full_version >= '3.14'", specifier = ">=3.11.1" }, { name = "dnspython", marker = "extra == 'srv'", specifier = ">=2.0.0" }, @@ -6264,13 +6269,14 @@ requires-dist = [ { name = "requests", specifier = ">=2.32.0,<3" }, { name = "requests-toolbelt", specifier = ">=1.0.0" }, ] -provides-extras = ["srv"] +provides-extras = ["srv", "standard"] [package.metadata.requires-dev] dev = [ { name = "apache-airflow", editable = "." }, { name = "apache-airflow-devel-common", editable = "devel-common" }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, + { name = "apache-airflow-providers-standard", editable = "providers/standard" }, { name = "apache-airflow-task-sdk", editable = "task-sdk" }, ] docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }]