ColtenOuO commented on code in PR #72133: URL: https://github.com/apache/airflow/pull/72133#discussion_r3869619614
########## providers/standard/src/airflow/providers/standard/sensors/websocket.py: ########## @@ -0,0 +1,92 @@ +# 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 +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 +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +if TYPE_CHECKING: + from airflow.sdk import Context + + +class WebSocketSensor(BaseSensorOperator): + """ + Waits for a message on a WebSocket connection. + + :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",) + + 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("Poking WebSocket %s", self.url) + with connect(self.url, additional_headers=self.header) as websocket: + if self.message_to_send is not None: + websocket.send(self.message_to_send) + try: + websocket.recv(timeout=self.poke_interval) + 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) + if not self.poke(context=context): Review Comment: In deferrable mode, this poke() first opens a WebSocket connection, sends message_to_send, and waits for one poke_interval. If the remote job hasn’t completed, poke() closes the connection, after which the trigger reconnects and sends the same request again. This can start the remote job twice and lose the response associated with the original connection. Since this feature is intended to wait on one long-lived WebSocket connection, the deferrable path should defer immediately and let the trigger own the only connection. ########## providers/standard/tests/unit/standard/triggers/test_websocket.py: ########## @@ -0,0 +1,77 @@ +# 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 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") + async def test_run_yields_event_with_received_message(self, mock_connect): + mock_websocket = mock.AsyncMock() Review Comment: Airflow’s testing guidelines require Mock, MagicMock, and AsyncMock instances to use spec or autospec, preventing tests from accepting methods or call signatures that the real client doesn’t provide. Both MagicMock instances here and the AsyncMock instances in the trigger tests currently have no spec. Please use the actual WebSocket connection type as the spec and consider applying autospec=True to mock.patch(). ########## providers/standard/tests/unit/standard/triggers/test_websocket.py: ########## @@ -0,0 +1,77 @@ +# 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 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") + async def test_run_yields_event_with_received_message(self, mock_connect): + mock_websocket = mock.AsyncMock() + 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") + async def test_run_does_not_send_without_message_to_send(self, mock_connect): + mock_websocket = mock.AsyncMock() Review Comment: ditto ########## uv.lock: ########## @@ -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" Review Comment: needs rebase (when provider release) ########## providers/standard/src/airflow/providers/standard/sensors/websocket.py: ########## @@ -0,0 +1,92 @@ +# 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 +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 +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +if TYPE_CHECKING: + from airflow.sdk import Context + + +class WebSocketSensor(BaseSensorOperator): Review Comment: This trigger resumes a deferred task, so it should inherit directly from BaseTrigger. BaseEventTrigger is the marker used for event-driven scheduling and asset watchers; on Airflow 3 it also sets supports_triggerer_queue=False, preventing this task-associated trigger from following task queue routing. Please use BaseTrigger on all supported Airflow versions, as FileTrigger does. ########## 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", Review Comment: The header states that this file is generated automatically and will be overwritten by the release workflow. The WebSocket sensor, trigger, and documentation metadata are already present in provider.yaml, so this file shouldn’t be synchronized manually; the next provider release will regenerate it from the source metadata. ########## providers/standard/src/airflow/providers/standard/sensors/websocket.py: ########## @@ -0,0 +1,92 @@ +# 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 +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 +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +if TYPE_CHECKING: + from airflow.sdk import Context + + +class WebSocketSensor(BaseSensorOperator): + """ + Waits for a message on a WebSocket connection. + + :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",) + + 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("Poking WebSocket %s", self.url) + with connect(self.url, additional_headers=self.header) as websocket: + if self.message_to_send is not None: + websocket.send(self.message_to_send) + try: + websocket.recv(timeout=self.poke_interval) + 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) Review Comment: super().execute() returns after consuming the first WebSocket message, but execution currently continues to the next line and calls self.poke() again. As a result, the non-deferrable sensor requires at least two messages to succeed; if the second receive times out, it can even call self.defer() while deferrable=False. The synchronous path must return immediately after it completes. More generally, because WebSocket messages are consumptive, the synchronous implementation should keep one connection open instead of repeatedly running the normal sensor poke loop. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
