dataops-sdk 0.4.1__tar.gz → 0.4.2__tar.gz

{dataops_sdk-0.4.1 → dataops_sdk-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataops-sdk
-Version: 0.4.1
+Version: 0.4.2
 Summary: Python SDK for the DataOps platform.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
@@ -21,6 +21,7 @@ Requires-Dist: doris-python>=1.0.3; extra == "runner"
 Requires-Dist: oracledb~=4.0.1; extra == "runner"
 Requires-Dist: neo4j; extra == "runner"
 Requires-Dist: kafka-python>3; extra == "runner"
+Requires-Dist: websockets>=13; extra == "runner"
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-mock; extra == "dev"

{dataops_sdk-0.4.1 → dataops_sdk-0.4.2}/pyproject.toml

@@ -34,6 +34,8 @@ runner = [
     "neo4j",
     # 队列驱动
     "kafka-python>3",
+    # WebSocket 客户端
+    "websockets>=13",
 ]
 dev = ["pytest", "pytest-mock", "ruff"]
 

{dataops_sdk-0.4.1 → dataops_sdk-0.4.2}/src/dataops_sdk/__init__.py

@@ -107,4 +107,4 @@ from .services import (  # noqa: F401
 )
 from .session import dataops_base_url_context, request_context
 
-__version__ = "0.4.1"
+__version__ = "0.4.2"

{dataops_sdk-0.4.1 → dataops_sdk-0.4.2}/src/dataops_sdk/runner/__init__.py

@@ -19,6 +19,7 @@ from .utils import (
     KafkaQueue,
     Neo4jConn,
     SFTPFileTransfer,
+    WebSocketConn,
     decrypt_body,
     encrypt_body,
     generate_uuid,
@@ -32,6 +33,7 @@ __all__ = [
     "FTPFileTransfer",
     "KafkaQueue",
     "SFTPFileTransfer",
+    "WebSocketConn",
     "Neo4jConn",
     "decrypt_body",
     "encrypt_body",

{dataops_sdk-0.4.1 → dataops_sdk-0.4.2}/src/dataops_sdk/runner/utils/__init__.py

@@ -33,6 +33,7 @@ from .nacos import get_service_instance
 from .neo4j_sync import Neo4jConn
 from .queue_kafka import KafkaQueue
 from .sftp_file_transfer import SFTPFileTransfer
+from .websocket_client import WebSocketConn
 
 __all__ = [
     "decrypt_body",
@@ -43,6 +44,7 @@ __all__ = [
     "Neo4jConn",
     "get_service_instance",
     "KafkaQueue",
+    "WebSocketConn",
     "FTPFileTransfer",
     "SFTPFileTransfer",
     "parse_args",

dataops_sdk-0.4.2/src/dataops_sdk/runner/utils/websocket_client.py

@@ -0,0 +1,240 @@
+"""WebSocket 客户端工具类。
+
+基于 websockets 的同步客户端实现，提供发送/接收数据（文本或二进制帧）能力。
+作为脚本内独立工具类使用，不接入平台数据源体系。
+
+同步 API 说明：
+- 构造即连接，连接失败抛出 ConnectionError；
+- receive 在超时或连接关闭时返回 None（对齐 KafkaQueue.recv 语义）；
+- 不做自动重连（单次脚本运行生命周期），留作后续扩展。
+
+注意：
+- receive(timeout=None) 会永久阻塞，脚本中拉取循环务必传入数值 timeout，
+  最终由执行器的子进程超时兜底；
+- max_size 默认 1MB，超限帧会抛出 ConnectionClosedError(1009) 并断连，
+  需要接收大帧时请调大 max_size；
+- receive_json 假定接收文本帧，二进制帧请使用 receive()。
+
+Example:
+    >>> with WebSocketConn(endpoint="example.com/ws", headers={"Authorization": "Bearer x"}, use_ssl=True) as ws:
+    ...     ws.send({"action": "subscribe", "channel": "trades"})
+    ...     while not ws.closed:
+    ...         msg = ws.receive(timeout=5)
+    ...         if msg is not None:
+    ...             print(msg)
+"""
+
+import json
+import os
+from logging import Logger
+from typing import Any
+
+from websockets import Subprotocol
+from websockets.exceptions import ConnectionClosed
+from websockets.protocol import State
+from websockets.sync.client import connect as _ws_connect
+
+from dataops_sdk.utils.log import getLogger
+
+__all__ = ["WebSocketConn"]
+
+
+class WebSocketConn:
+    """同步 WebSocket 客户端连接。
+
+    构造即建立连接，提供 send / receive / send_json / receive_json 方法，
+    支持上下文管理器（退出时自动关闭）。
+
+    Attributes:
+        logger: 日志记录器，用于记录连接操作日志。
+        url: WebSocket 服务端地址。
+
+    Example:
+        >>> ws = WebSocketConn(endpoint="example.com/ws", use_ssl=True)
+        >>> ws.send({"action": "subscribe"})
+        >>> msg = ws.receive(timeout=5)
+        >>> ws.close()
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        *,
+        use_ssl: bool = False,
+        headers: dict[str, str] | None = None,
+        protocols: list[Subprotocol] | None = None,
+        open_timeout: float | None = 10.0,
+        ping_interval: float | None = 20.0,
+        ping_timeout: float | None = 20.0,
+        close_timeout: float | None = 10.0,
+        max_size: int | None = 2**20,
+        task_id: str | None = None,
+    ):
+        """初始化并建立 WebSocket 连接。
+
+        Args:
+            endpoint: WebSocket 服务端地址，只需输入主机:端口（可带路径），
+                如 "host:9001" 或 "host:9001/ws"；若已含协议前缀则按原样使用。
+            use_ssl: 是否使用 TLS（wss://），默认 False 使用 ws://。
+            headers: 附加请求头（如 Authorization），可选。
+            protocols: 子协议列表，可选。
+            open_timeout: 建立连接超时（秒）。
+            ping_interval: 心跳间隔（秒）。
+            ping_timeout: 心跳超时（秒）。
+            close_timeout: 关闭连接超时（秒）。
+            max_size: 单帧最大字节数，超限帧将被关闭连接。
+            task_id: 任务 ID，用于日志标识；缺省时使用 TASK_ID 环境变量。
+
+        Raises:
+            ConnectionError: 连接失败时抛出。
+        """
+        self._endpoint = endpoint
+        if "://" in endpoint:
+            self._url = endpoint
+        else:
+            scheme = "wss" if use_ssl else "ws"
+            self._url = f"{scheme}://{endpoint}"
+        # 先初始化为 None，避免构造失败时 __del__ 调用 close 访问不存在的属性
+        self._conn = None
+        self._logger = getLogger(task_id or os.getenv("TASK_ID") or type(self).__name__)
+        try:
+            self._conn = _ws_connect(
+                self._url,
+                subprotocols=protocols,
+                additional_headers=headers,
+                open_timeout=open_timeout,
+                ping_interval=ping_interval,
+                ping_timeout=ping_timeout,
+                close_timeout=close_timeout,
+                max_size=max_size,
+            )
+        except Exception as e:
+            self._logger.error(f"WebSocket 连接失败: {self._url}: {e}")
+            raise ConnectionError(f"WebSocket connect failed: {self._url}: {e}") from e
+        self._logger.info(f"WebSocket 已连接: {self._url}")
+
+    def send(self, data: str | bytes | dict | list) -> None:
+        """发送消息。
+
+        str 发送为文本帧，bytes 发送为二进制帧；
+        dict / list 自动序列化为 JSON 文本帧。
+
+        Args:
+            data: 待发送数据。
+
+        Raises:
+            ConnectionError: 连接已关闭时抛出。
+        """
+        if self._conn is None or self._conn.state is State.CLOSED:
+            raise ConnectionError(f"WebSocket 连接已关闭，无法发送数据: {self._url}")
+        if isinstance(data, (dict, list)):
+            data = json.dumps(data, ensure_ascii=False)
+        try:
+            self._conn.send(data)
+        except ConnectionClosed as e:
+            self._logger.error(f"WebSocket 发送失败（连接已关闭）: {e}")
+            raise
+        except Exception as e:
+            self._logger.error(f"WebSocket 发送失败: {e}")
+            raise
+
+    def receive(self, timeout: float | None = None) -> str | bytes | None:
+        """接收一条消息。
+
+        Args:
+            timeout: 等待超时（秒），None 表示阻塞等待。
+
+        Returns:
+            str | bytes | None: 文本帧返回 str，二进制帧返回 bytes；
+            超时或连接关闭时返回 None。
+        """
+        if self._conn is None or self._conn.state is State.CLOSED:
+            return None
+        try:
+            return self._conn.recv(timeout=timeout)
+        except TimeoutError:
+            return None
+        except ConnectionClosed as e:
+            self._logger.warning(f"WebSocket 连接已关闭: {e}")
+            return None
+
+    def send_json(self, obj: Any) -> None:
+        """发送 JSON 消息。
+
+        Args:
+            obj: 任意可 JSON 序列化的对象。
+        """
+        self.send(json.dumps(obj, ensure_ascii=False))
+
+    def receive_json(self, timeout: float | None = None) -> Any | None:
+        """接收并解析一条 JSON 文本帧。
+
+        Args:
+            timeout: 等待超时（秒），None 表示阻塞等待。
+
+        Returns:
+            Any | None: 解析后的对象；超时或连接关闭时返回 None。
+        """
+        raw = self.receive(timeout=timeout)
+        if raw is None:
+            return None
+        return json.loads(raw)
+
+    def close(self) -> None:
+        """关闭连接。
+
+        幂等，可多次调用。
+        """
+        if self._conn is None:
+            return
+        try:
+            if self._conn.state is not State.CLOSED:
+                self._conn.close()
+            self._logger.info("WebSocket 已关闭")
+        except Exception as e:
+            self._logger.error(f"WebSocket 关闭失败: {e}")
+        finally:
+            self._conn = None
+
+    @property
+    def logger(self) -> Logger:
+        """日志记录器。"""
+        return self._logger
+
+    @property
+    def url(self) -> str:
+        """完整的 WebSocket 地址（含协议前缀）。"""
+        return self._url
+
+    @property
+    def endpoint(self) -> str:
+        """WebSocket 服务端地址（不含协议前缀）。"""
+        return self._endpoint
+
+    @property
+    def closed(self) -> bool:
+        """连接是否已关闭。"""
+        return self._conn is None or self._conn.state is State.CLOSED
+
+    @property
+    def connected(self) -> bool:
+        """连接是否处于打开状态。"""
+        return self._conn is not None and self._conn.state is State.OPEN
+
+    def __enter__(self) -> "WebSocketConn":
+        """上下文管理器入口。
+
+        Returns:
+            WebSocketConn: 当前实例。
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """上下文管理器退出，自动关闭连接。"""
+        if exc_type:
+            self._logger.exception(exc_val)
+        self.close()
+
+    def __del__(self):
+        """析构时兜底关闭连接。"""
+        self.close()