diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index b0cd8ce92ea2c..c91c0cbc45b7b 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -454,6 +454,7 @@ decrypted Decrypts deduplicate deduplicated +deduplicating deduplication deepcopy DefaultAzureCredential diff --git a/providers/standard/docs/index.rst b/providers/standard/docs/index.rst index a6b8f12bab228..ffd6d9a48dc40 100644 --- a/providers/standard/docs/index.rst +++ b/providers/standard/docs/index.rst @@ -127,6 +127,7 @@ Install them when installing from PyPI. For example: Extra Dependencies =============== ======================================== ``openlineage`` ``apache-airflow-providers-openlineage`` +``websocket`` ``websockets>=14.0`` =============== ======================================== Downloading official packages diff --git a/providers/standard/docs/sensors/websocket.rst b/providers/standard/docs/sensors/websocket.rst new file mode 100644 index 0000000000000..72000b387fdda --- /dev/null +++ b/providers/standard/docs/sensors/websocket.rst @@ -0,0 +1,61 @@ + .. 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. + + + +.. _howto/operator:WebSocketSensor: + +WebSocketSensor +================ + +Use the :class:`~airflow.providers.standard.sensors.websocket.WebSocketSensor` to wait for a message on a +``ws://`` or ``wss://`` WebSocket connection. This is useful when a remote server accepts a long-lived +connection and replies asynchronously once it has handled the requested work, since a plain HTTP request +is not well suited to that kind of long-lived connection. Requires the ``websocket`` extra +(``apache-airflow-providers-standard[websocket]``). + +.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_sensors.py + :language: python + :dedent: 4 + :start-after: [START example_websocket_sensor] + :end-before: [END example_websocket_sensor] + +Also for this job you can use sensor in the deferrable mode: + +.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_sensors.py + :language: python + :dedent: 4 + :start-after: [START example_websocket_sensor_async] + :end-before: [END example_websocket_sensor_async] + +A common use case is to send a request over the connection right after it opens (via +``message_to_send``) and then wait for the remote server's asynchronous reply, optionally +passing connection headers such as an auth token via ``header``: + +.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_sensors.py + :language: python + :dedent: 4 + :start-after: [START example_websocket_sensor_send_message_async] + :end-before: [END example_websocket_sensor_send_message_async] + +.. warning:: + In deferrable mode, ``message_to_send`` may be sent more than once for a single task + run. Airflow triggers are not guaranteed to execute exactly once — a triggerer + restart or redistribution to another host re-runs the trigger from scratch, opening + a new connection and re-sending ``message_to_send``. If that message has a side + effect on the remote server, such as starting a job, the server must treat a resend + as safe — for example by deduplicating on a request id embedded in the message. diff --git a/providers/standard/provider.yaml b/providers/standard/provider.yaml index d593430a0ffd2..40c39ad3caea1 100644 --- a/providers/standard/provider.yaml +++ b/providers/standard/provider.yaml @@ -83,6 +83,7 @@ integrations: - /docs/apache-airflow-providers-standard/sensors/datetime.rst - /docs/apache-airflow-providers-standard/sensors/file.rst - /docs/apache-airflow-providers-standard/sensors/external_task_sensor.rst + - /docs/apache-airflow-providers-standard/sensors/websocket.rst operators: - integration-name: Standard @@ -108,6 +109,7 @@ sensors: - airflow.providers.standard.sensors.python - airflow.providers.standard.sensors.filesystem - airflow.providers.standard.sensors.external_task + - airflow.providers.standard.sensors.websocket hooks: - integration-name: Standard python-modules: @@ -122,6 +124,7 @@ triggers: - airflow.providers.standard.triggers.file - airflow.providers.standard.triggers.temporal - airflow.providers.standard.triggers.hitl + - airflow.providers.standard.triggers.websocket extra-links: - airflow.providers.standard.operators.trigger_dagrun.TriggerDagRunLink diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index 04e8c68d8c91a..c19b23befef29 100644 --- a/providers/standard/pyproject.toml +++ b/providers/standard/pyproject.toml @@ -69,6 +69,9 @@ dependencies = [ "openlineage" = [ "apache-airflow-providers-openlineage" ] +"websocket" = [ + "websockets>=14.0" +] [dependency-groups] dev = [ @@ -79,6 +82,7 @@ dev = [ "apache-airflow-providers-openlineage", # Additional devel dependencies (do not remove this line and add extra development dependencies) "apache-airflow-providers-mysql", + "websockets>=14.0", ] # To build docs: diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py b/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py index 73ccaadbe0f02..a1fcdc64538c6 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py @@ -28,6 +28,7 @@ from airflow.providers.standard.sensors.python import PythonSensor from airflow.providers.standard.sensors.time import TimeSensor from airflow.providers.standard.sensors.time_delta import TimeDeltaSensor +from airflow.providers.standard.sensors.websocket import WebSocketSensor from airflow.providers.standard.sensors.weekday import DayOfWeekSensor from airflow.providers.standard.utils.weekday import WeekDay from airflow.sdk import DAG @@ -125,6 +126,34 @@ def failure_callable(): ) # [END example_day_of_week_sensor] + # [START example_websocket_sensor] + t12 = WebSocketSensor( + task_id="wait_for_websocket_message", url="wss://example.com/socket", timeout=3, soft_fail=True + ) + # [END example_websocket_sensor] + + # [START example_websocket_sensor_async] + t13 = WebSocketSensor( + task_id="wait_for_websocket_message_async", + url="wss://example.com/socket", + deferrable=True, + timeout=3, + soft_fail=True, + ) + # [END example_websocket_sensor_async] + + # [START example_websocket_sensor_send_message_async] + t14 = WebSocketSensor( + task_id="request_and_wait_for_websocket_reply", + url="wss://example.com/socket", + message_to_send='{"action": "start_job"}', + header={"Authorization": "Bearer my-token"}, + deferrable=True, + timeout=3, + soft_fail=True, + ) + # [END example_websocket_sensor_send_message_async] + tx = BashOperator(task_id="print_date_in_bash", bash_command="date") tx.trigger_rule = TriggerRule.NONE_FAILED @@ -133,3 +162,4 @@ def failure_callable(): t8 >> tx [t9, t10] >> tx t11 >> tx + [t12, t13, t14] >> tx diff --git a/providers/standard/src/airflow/providers/standard/get_provider_info.py b/providers/standard/src/airflow/providers/standard/get_provider_info.py index 1f7b2049454d1..0be729079a0cb 100644 --- a/providers/standard/src/airflow/providers/standard/get_provider_info.py +++ b/providers/standard/src/airflow/providers/standard/get_provider_info.py @@ -43,6 +43,7 @@ def get_provider_info(): "/docs/apache-airflow-providers-standard/sensors/datetime.rst", "/docs/apache-airflow-providers-standard/sensors/file.rst", "/docs/apache-airflow-providers-standard/sensors/external_task_sensor.rst", + "/docs/apache-airflow-providers-standard/sensors/websocket.rst", ], } ], @@ -75,6 +76,7 @@ def get_provider_info(): "airflow.providers.standard.sensors.python", "airflow.providers.standard.sensors.filesystem", "airflow.providers.standard.sensors.external_task", + "airflow.providers.standard.sensors.websocket", ], } ], @@ -96,6 +98,7 @@ def get_provider_info(): "airflow.providers.standard.triggers.file", "airflow.providers.standard.triggers.temporal", "airflow.providers.standard.triggers.hitl", + "airflow.providers.standard.triggers.websocket", ], } ], diff --git a/providers/standard/src/airflow/providers/standard/sensors/websocket.py b/providers/standard/src/airflow/providers/standard/sensors/websocket.py new file mode 100644 index 0000000000000..bf87bffd36af9 --- /dev/null +++ b/providers/standard/src/airflow/providers/standard/sensors/websocket.py @@ -0,0 +1,107 @@ +# 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 + +import datetime +import time +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +from websockets.sync.client import connect + +from airflow.providers.common.compat.sdk import BaseSensorOperator, conf, poke_mode_only +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +if TYPE_CHECKING: + from airflow.sdk import Context + + +@poke_mode_only +class WebSocketSensor(BaseSensorOperator): + """ + Waits for a message on a WebSocket connection. + + WebSocket messages are consumptive: once read, a message cannot be read again, and + reconnecting can duplicate ``message_to_send`` against the remote server. The + non-deferrable path therefore opens exactly one connection and blocks on it for up to + ``timeout`` seconds instead of reconnecting every ``poke_interval`` (which has no + effect on this sensor), and this sensor is marked poke-mode-only since a rescheduled + invocation would need a new connection anyway. + + :param url: The ``ws://`` or ``wss://`` URL of the WebSocket server to connect to. + :param header: Optional headers sent when opening the connection. + :param message_to_send: Optional message sent right after the connection is established. + :param deferrable: If waiting for completion, whether to defer the task until done, + default is ``False``. + + .. seealso:: + For more information on how to use this sensor, take a look at the guide: + :ref:`howto/operator:WebSocketSensor` + """ + + template_fields: Sequence[str] = ("url", "header", "message_to_send") + template_fields_renderers = {"header": "json"} + + def __init__( + self, + *, + url: str, + header: dict[str, str] | None = None, + message_to_send: str | bytes | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + **kwargs, + ): + super().__init__(**kwargs) + self.url = url + self.header = header + self.message_to_send = message_to_send + self.deferrable = deferrable + + def poke(self, context: Context) -> bool: + self.log.info("Connecting to WebSocket %s", self.url) + deadline = time.monotonic() + self.timeout + try: + with connect(self.url, additional_headers=self.header, open_timeout=self.timeout) as websocket: + if self.message_to_send is not None: + websocket.send(self.message_to_send) + websocket.recv(timeout=max(deadline - time.monotonic(), 0)) + except TimeoutError: + return False + self.log.info("Received message from %s", self.url) + return True + + def execute(self, context: Context) -> None: + if not self.deferrable: + super().execute(context=context) + return + # Each poke opens and consumes a WebSocket connection, so the deferrable path must + # defer immediately: polling here first would send message_to_send and consume the + # reply before handing off, leaving the trigger to open a second connection and + # re-send the request. + self.defer( + timeout=datetime.timedelta(seconds=self.timeout), + trigger=WebSocketTrigger( + url=self.url, + header=self.header, + message_to_send=self.message_to_send, + ), + method_name="execute_complete", + ) + + def execute_complete(self, context: Context, event: Any = None) -> None: + """Handle the event when the trigger fires and return immediately.""" + self.log.info("%s completed successfully with message: %s", self.task_id, event) diff --git a/providers/standard/src/airflow/providers/standard/triggers/websocket.py b/providers/standard/src/airflow/providers/standard/triggers/websocket.py new file mode 100644 index 0000000000000..c3436a5b8ac0c --- /dev/null +++ b/providers/standard/src/airflow/providers/standard/triggers/websocket.py @@ -0,0 +1,76 @@ +# 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 AsyncIterator +from typing import Any + +from websockets.asyncio.client import connect + +from airflow.triggers.base import BaseTrigger, TriggerEvent + + +class WebSocketTrigger(BaseTrigger): + """ + A trigger that opens a WebSocket connection and fires once a message is received. + + This is meant for deferrable operators that hand off a long-lived request to a remote + WebSocket server and resume once that server replies, without occupying a worker slot + while waiting. + + Like any Airflow trigger, ``run()`` is not guaranteed to execute only once: a + triggerer restart or redistribution to another host re-runs it from scratch. Each + execution opens a new connection and re-sends ``message_to_send`` if one is set, so + if that message starts a remote job, the remote server must treat a resend as safe — + for example by deduplicating on a request id embedded in the message. + + :param url: The ``ws://`` or ``wss://`` URL of the WebSocket server to connect to. + :param header: Optional headers sent when opening the connection. + :param message_to_send: Optional message sent right after the connection is established. + """ + + def __init__( + self, + url: str, + header: dict[str, str] | None = None, + message_to_send: str | bytes | None = None, + **kwargs, + ): + super().__init__() + self.url = url + self.header = header + self.message_to_send = message_to_send + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize WebSocketTrigger arguments and classpath.""" + return ( + "airflow.providers.standard.triggers.websocket.WebSocketTrigger", + { + "url": self.url, + "header": self.header, + "message_to_send": self.message_to_send, + }, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Connect to the WebSocket server and wait for the first message.""" + async with connect(self.url, additional_headers=self.header) as websocket: + if self.message_to_send is not None: + await websocket.send(self.message_to_send) + message = await websocket.recv() + self.log.info("Received message from %s", self.url) + yield TriggerEvent(message) diff --git a/providers/standard/tests/unit/standard/sensors/test_websocket.py b/providers/standard/tests/unit/standard/sensors/test_websocket.py new file mode 100644 index 0000000000000..45ef4e8435036 --- /dev/null +++ b/providers/standard/tests/unit/standard/sensors/test_websocket.py @@ -0,0 +1,148 @@ +# 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 unittest import mock + +import pytest +from websockets.sync.client import ClientConnection + +from airflow.models.dag import DAG +from airflow.providers.common.compat.sdk import AirflowSensorTimeout, TaskDeferred +from airflow.providers.standard.sensors.websocket import WebSocketSensor +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +from tests_common.test_utils.version_compat import timezone + +URL = "ws://example.com/socket" +DEFAULT_DATE = timezone.datetime(2015, 1, 1) + + +class TestWebSocketSensor: + @classmethod + def setup_class(cls): + args = {"owner": "airflow", "start_date": DEFAULT_DATE} + cls.dag = DAG("test_websocket_sensor", schedule=None, default_args=args) + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_returns_true_on_message(self, mock_connect): + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor(task_id="poke_true", url=URL, message_to_send="ping", dag=self.dag) + assert sensor.poke(context={}) is True + mock_websocket.send.assert_called_once_with("ping") + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_returns_false_on_timeout(self, mock_connect): + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.side_effect = TimeoutError() + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor(task_id="poke_false", url=URL, dag=self.dag) + assert sensor.poke(context={}) is False + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_waits_for_the_overall_timeout_not_poke_interval(self, mock_connect): + """recv() must be bounded by the sensor's overall timeout, not poke_interval — + otherwise a single poke could block past the sensor's declared timeout before + that timeout is ever checked.""" + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor( + task_id="poke_timeout_arg", url=URL, timeout=45, poke_interval=5, dag=self.dag + ) + assert sensor.poke(context={}) is True + + mock_connect.assert_called_once_with(URL, additional_headers=None, open_timeout=45) + (_, kwargs) = mock_websocket.recv.call_args + assert kwargs["timeout"] == pytest.approx(45, abs=1) + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_returns_false_when_handshake_times_out(self, mock_connect): + """A connect() timeout (slow handshake) must be treated the same as a recv() + timeout — including respecting soft_fail — not propagate as a raw TimeoutError + that bypasses the sensor's normal timeout handling.""" + mock_connect.side_effect = TimeoutError("timed out while waiting for handshake response") + + sensor = WebSocketSensor(task_id="handshake_timeout", url=URL, dag=self.dag) + assert sensor.poke(context={}) is False + + @mock.patch("airflow.providers.standard.sensors.websocket.time.monotonic") + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_recv_gets_remaining_time_after_slow_handshake(self, mock_connect, mock_monotonic): + """If the handshake itself consumes part of the timeout budget, recv() must only + get what's left, not the full timeout again — otherwise total wait time could + exceed the sensor's declared timeout.""" + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value.__enter__.return_value = mock_websocket + # deadline computed at t=0 with timeout=10; connect() "takes" 4s, leaving 6s for recv(). + mock_monotonic.side_effect = [0, 4] + + sensor = WebSocketSensor(task_id="slow_handshake", url=URL, timeout=10, dag=self.dag) + assert sensor.poke(context={}) is True + mock_websocket.recv.assert_called_once_with(timeout=6) + + def test_reschedule_mode_not_allowed(self): + with pytest.raises(ValueError, match="Cannot set mode to 'reschedule'. Only 'poke' is acceptable"): + WebSocketSensor(task_id="reschedule", url=URL, mode="reschedule", dag=self.dag) + + def test_task_defer_does_not_poke_first(self): + """The deferrable path must defer immediately: poke() consumes the connection, + so polling before deferring would send message_to_send and lose the reply the + trigger is supposed to wait for.""" + sensor = WebSocketSensor(task_id="defer", url=URL, deferrable=True, dag=self.dag) + + with mock.patch.object(WebSocketSensor, "poke", autospec=True) as mock_poke: + with pytest.raises(TaskDeferred) as exc: + sensor.execute({}) + + mock_poke.assert_not_called() + assert isinstance(exc.value.trigger, WebSocketTrigger) + assert exc.value.trigger.url == URL + + def test_execute_sync_calls_poke_exactly_once(self): + """Since poke() already blocks for the full sensor timeout, execute() must never + call it a second time — a second call would open a new connection and re-send + message_to_send.""" + sensor = WebSocketSensor(task_id="sync_timeout", url=URL, timeout=0, dag=self.dag) + + with mock.patch.object(WebSocketSensor, "poke", autospec=True, return_value=False) as mock_poke: + with pytest.raises(AirflowSensorTimeout): + sensor.execute({}) + + mock_poke.assert_called_once() + + def test_template_fields_are_rendered(self): + """url, header, and message_to_send commonly need runtime values (run_id, an + idempotency key, an auth token), so all three must be templated.""" + sensor = WebSocketSensor( + task_id="templated", + url="wss://example.com/{{ run_id }}", + message_to_send='{"run_id": "{{ run_id }}"}', + header={"Authorization": "Bearer {{ run_id }}"}, + dag=self.dag, + ) + sensor.render_template_fields({"run_id": "manual__2024-01-01"}) + + assert sensor.url == "wss://example.com/manual__2024-01-01" + assert sensor.message_to_send == '{"run_id": "manual__2024-01-01"}' + assert sensor.header == {"Authorization": "Bearer manual__2024-01-01"} diff --git a/providers/standard/tests/unit/standard/triggers/test_websocket.py b/providers/standard/tests/unit/standard/triggers/test_websocket.py new file mode 100644 index 0000000000000..15cdb767cfbe7 --- /dev/null +++ b/providers/standard/tests/unit/standard/triggers/test_websocket.py @@ -0,0 +1,99 @@ +# 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 unittest import mock + +import pytest +from websockets.asyncio.client import ClientConnection + +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + + +class _FakeConnection: + """Stands in for the object returned by ``websockets.asyncio.client.connect``.""" + + def __init__(self, websocket): + self._websocket = websocket + + async def __aenter__(self): + return self._websocket + + async def __aexit__(self, *exc_info): + return False + + +class TestWebSocketTrigger: + URL = "ws://example.com/socket" + + def test_serialization(self): + """Asserts that the trigger correctly serializes its arguments and classpath.""" + trigger = WebSocketTrigger(url=self.URL, header={"Authorization": "token"}, message_to_send="ping") + classpath, kwargs = trigger.serialize() + assert classpath == "airflow.providers.standard.triggers.websocket.WebSocketTrigger" + assert kwargs == { + "url": self.URL, + "header": {"Authorization": "token"}, + "message_to_send": "ping", + } + + @pytest.mark.asyncio + @mock.patch("airflow.providers.standard.triggers.websocket.connect", autospec=True) + async def test_run_yields_event_with_received_message(self, mock_connect): + mock_websocket = mock.AsyncMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value = _FakeConnection(mock_websocket) + + trigger = WebSocketTrigger(url=self.URL, header={"Authorization": "token"}, message_to_send="ping") + event = await trigger.run().__anext__() + + mock_connect.assert_called_once_with(self.URL, additional_headers={"Authorization": "token"}) + mock_websocket.send.assert_awaited_once_with("ping") + assert event.payload == "pong" + + @pytest.mark.asyncio + @mock.patch("airflow.providers.standard.triggers.websocket.connect", autospec=True) + async def test_run_does_not_send_without_message_to_send(self, mock_connect): + mock_websocket = mock.AsyncMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value = _FakeConnection(mock_websocket) + + trigger = WebSocketTrigger(url=self.URL) + await trigger.run().__anext__() + + mock_websocket.send.assert_not_awaited() + + @pytest.mark.asyncio + @mock.patch("airflow.providers.standard.triggers.websocket.connect", autospec=True) + async def test_reconstructed_trigger_resends_message_to_send(self, mock_connect): + """Documents that a triggerer restart or redistribution — which reconstructs the + trigger from its serialize() output and calls run() again — re-sends + message_to_send. Airflow does not guarantee a trigger's run() executes only + once, so callers whose message starts a remote job must make that job + idempotent; this is not something the trigger itself can enforce.""" + mock_websocket = mock.AsyncMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value = _FakeConnection(mock_websocket) + + original = WebSocketTrigger(url=self.URL, message_to_send="start_job") + _, kwargs = original.serialize() + reconstructed = WebSocketTrigger(**kwargs) + + await original.run().__anext__() + await reconstructed.run().__anext__() + + assert mock_websocket.send.await_args_list == [mock.call("start_job"), mock.call("start_job")] diff --git a/uv.lock b/uv.lock index d5a35a4b581eb..09a2b237cba00 100644 --- a/uv.lock +++ b/uv.lock @@ -3205,7 +3205,7 @@ docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "d [[package]] name = "apache-airflow-providers-amazon" -version = "9.35.0" +version = "9.35.1" source = { editable = "providers/amazon" } dependencies = [ { name = "apache-airflow" }, @@ -6587,7 +6587,7 @@ docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "d [[package]] name = "apache-airflow-providers-microsoft-azure" -version = "15.0.0" +version = "15.0.1" source = { editable = "providers/microsoft/azure" } dependencies = [ { name = "adlfs" }, @@ -8272,6 +8272,9 @@ dependencies = [ openlineage = [ { name = "apache-airflow-providers-openlineage" }, ] +websocket = [ + { name = "websockets" }, +] [package.dev-dependencies] dev = [ @@ -8281,6 +8284,7 @@ dev = [ { name = "apache-airflow-providers-mysql" }, { name = "apache-airflow-providers-openlineage" }, { name = "apache-airflow-task-sdk" }, + { name = "websockets" }, ] docs = [ { name = "apache-airflow-devel-common", extra = ["docs"] }, @@ -8291,8 +8295,9 @@ requires-dist = [ { name = "apache-airflow", editable = "." }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-openlineage", marker = "extra == 'openlineage'", editable = "providers/openlineage" }, + { name = "websockets", marker = "extra == 'websocket'", specifier = ">=14.0" }, ] -provides-extras = ["openlineage"] +provides-extras = ["openlineage", "websocket"] [package.metadata.requires-dev] dev = [ @@ -8302,6 +8307,7 @@ dev = [ { name = "apache-airflow-providers-mysql", editable = "providers/mysql" }, { name = "apache-airflow-providers-openlineage", editable = "providers/openlineage" }, { name = "apache-airflow-task-sdk", editable = "task-sdk" }, + { name = "websockets", specifier = ">=14.0" }, ] docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }]