Add WebSocketSensor and WebSocketTrigger to the standard provider

docs/spelling_wordlist.txt

@@ -454,6 +454,7 @@ decrypted
 Decrypts
 deduplicate
 deduplicated
+deduplicating
 deduplication
 deepcopy
 DefaultAzureCredential

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

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.

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

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:

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

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

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)

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)