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" }]