ColtenOuO commented on code in PR #72133: URL: https://github.com/apache/airflow/pull/72133#discussion_r3871349363
########## providers/standard/src/airflow/providers/standard/sensors/websocket.py: ########## @@ -0,0 +1,112 @@ +# 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 ClientConnection, 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 keeps a single connection open across pokes instead of + reconnecting and re-sending on every poke, and this sensor will not behave correctly in + reschedule mode, since that state would be lost between rescheduled invocations. + + :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 + self._connection: ClientConnection | None = None Review Comment: Although the extra poke() after a successful synchronous execution has been fixed, this receive still times out and closes the connection whenever the server takes longer than one poke_interval. BaseSensorOperator.execute() then calls poke() again, reconnecting and resending message_to_send. A long-running remote job can therefore still be started multiple times, while the response associated with the original connection is lost. Also, with timeout=3 and the default poke_interval=60, a single recv() can block for 60 seconds and exceed the sensor’s overall timeout. The synchronous path should wait on one connection until the sensor timeout rather than use a reconnecting poke loop; reschedule mode should also be explicitly rejected or handled. ########## providers/standard/src/airflow/providers/standard/sensors/websocket.py: ########## @@ -0,0 +1,112 @@ +# 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 ClientConnection, 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 keeps a single connection open across pokes instead of + reconnecting and re-sending on every poke, and this sensor will not behave correctly in + reschedule mode, since that state would be lost between rescheduled invocations. + + :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",) Review Comment: Only url is currently templated, but real requests commonly need message_to_send to contain the run_id, task-instance information, or a stable idempotency key, while authorization headers often need runtime resolution. The current API pushes users toward hard-coding these values in Dag source, as the documentation example does with its token. Please add header and message_to_send to template_fields, add a JSON renderer for the headers, and cover their rendering in tests. ########## providers/standard/src/airflow/providers/standard/triggers/websocket.py: ########## @@ -0,0 +1,70 @@ +# 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. + + :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) Review Comment: Airflow triggers must assume that run() may execute more than once, for example after triggerer restart, redistribution, or a network partition. Each execution here sends message_to_send again, while the documentation example uses {"action": "start_job"}, so one task execution may start multiple remote jobs. Airflow can deduplicate duplicate TriggerEvent objects, but it cannot undo external side effects that have already occurred. The protocol needs a stable idempotency or correlation key, or it must explicitly require the server to treat retransmissions as the same request; otherwise, the side effect that starts the job should be moved out of the trigger. Please also add documentation and a test covering trigger reconstruction. -- 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]
