dataops-sdk 0.3.0__py3-none-any.whl

dataops_sdk/runner/utils/file_transfer_base.py

@@ -0,0 +1,307 @@
+"""
+文件传输抽象基类模块。
+
+提供统一的 FTP/SFTP 文件传输接口，支持上传、下载、列表等操作。
+
+Classes:
+    FileTransferBase: 文件传输操作抽象基类
+
+Example:
+    >>> class MyFileTransfer(FileTransferBase):
+    ...     def post_init(self):
+    ...         pass  # 初始化具体传输连接
+    ...     def upload(self, local_path, remote_name):
+    ...         pass
+    ...     def download(self, remote_name, local_path):
+    ...         pass
+    ...     def list_files(self, pattern=None):
+    ...         return []
+    ...     def delete(self, remote_name):
+    ...         pass
+    ...     def exists(self, remote_name):
+    ...         pass
+    ...     def ping(self):
+    ...         pass
+    ...     def close(self):
+    ...         pass
+    ...
+    >>> with MyFileTransfer(host="localhost", port=21,
+    ...                    user="user", password="pass") as ft:
+    ...     ft.upload("/local/file.txt", "remote_file.txt")
+"""
+
+import os
+from abc import ABC, abstractmethod
+from logging import Logger
+from typing import BinaryIO
+
+from .log import getLogger
+
+
+class FileTransferBase(ABC):
+    """
+    文件传输抽象基类。
+
+    定义文件传输的标准接口，包括上传、下载、列表、删除和连接检查等方法。
+
+    继承此类需要实现以下抽象方法：
+    - post_init: 初始化具体传输连接对象
+    - upload: 上传文件到远程服务器
+    - download: 从远程服务器下载文件
+    - list_files: 列出远程文件
+    - delete: 删除远程文件
+    - exists: 检查远程文件是否存在
+    - ping: 检查连接是否存活
+    - close: 关闭连接
+
+    可以作为上下文管理器使用，会自动在退出时关闭连接。
+
+    Attributes:
+        logger: 日志记录器，用于记录文件传输操作日志。
+
+    Example:
+        >>> with FTPFileTransfer(host="ftp.example.com", port=21,
+        ...                     user="user", password="pass") as ftp:
+        ...     ftp.upload("/local/file.txt", "remote_file.txt")
+        ...     files = ftp.list_files("*.txt")
+    """
+
+    def __init__(
+        self,
+        host: str,
+        user: str,
+        password: str,
+        remote_path: str = "/",
+        port: int | str = 22,
+        name: str | None = None,
+    ):
+        """
+        初始化文件传输基类。
+
+        将传入的连接参数保存为实例属性，并初始化日志记录器。
+        子类应在 post_init 方法中创建实际的文件传输连接。
+
+        Args:
+            host: 文件服务器主机地址。
+            port: 文件服务器端口号。
+            user: 文件服务器用户名。
+            password: 文件服务器密码。
+            remote_path: 远程基础目录，默认为 "/"。
+            name: 连接名称。
+        """
+        self._host = host
+        self._port = int(port)
+        self._user = user
+        self._password = password
+        self._remote_path = remote_path
+
+        self._name = name or f"file_transfer_{self._host}"
+
+        self._logger = getLogger(name)
+
+        self.post_init()
+
+    @abstractmethod
+    def post_init(self):
+        """
+        初始化具体文件传输连接对象。
+
+        子类应在此方法中创建实际的 FTP/SFTP 连接。
+        该方法在 __init__ 中被调用，用于延迟初始化连接。
+
+        Returns:
+            None
+        """
+        pass
+
+    @abstractmethod
+    def upload(self, fp: BinaryIO | str | bytes, remote_name: str) -> bool:
+        """
+        上传文件到远程服务器。
+
+        将本地文件对象上传到远程服务器，存储在 remote_path 目录下。
+
+        Args:
+            fp: 文件对象，支持任意 IO 类型（如 BytesIO, FileIO 等）。
+            remote_name: 远程文件名（不含路径）。
+
+        Returns:
+            bool: 上传成功返回 True，否则返回 False。
+        """
+        pass
+
+    @abstractmethod
+    def download(self, remote_name: str, local_path: str) -> bool:
+        """
+        从远程服务器下载文件。
+
+        从远程服务器下载文件到本地路径。
+
+        Args:
+            remote_name: 远程文件名（不含路径）。
+            local_path: 本地保存路径（包含文件名）。
+
+        Returns:
+            bool: 下载成功返回 True，否则返回 False。
+        """
+        pass
+
+    @abstractmethod
+    def list_files(self, pattern: str | None = None) -> list[str]:
+        """
+        列出远程目录下的文件。
+
+        Args:
+            pattern: 可选的 glob 模式，用于过滤文件名。
+
+        Returns:
+            list[str]: 文件名列表。
+        """
+        pass
+
+    @abstractmethod
+    def delete(self, remote_name: str) -> bool:
+        """
+        删除远程文件。
+
+        Args:
+            remote_name: 远程文件名（不含路径）。
+
+        Returns:
+            bool: 删除成功返回 True，否则返回 False。
+        """
+        pass
+
+    @abstractmethod
+    def exists(self, remote_name: str) -> bool:
+        """
+        检查远程文件是否存在。
+
+        Args:
+            remote_name: 远程文件名（不含路径）。
+
+        Returns:
+            bool: 文件存在返回 True，否则返回 False。
+        """
+        pass
+
+    @abstractmethod
+    def ping(self) -> bool:
+        """
+        检查连接是否存活。
+
+        Returns:
+            bool: 连接正常返回 True，否则返回 False。
+        """
+        pass
+
+    @abstractmethod
+    def close(self):
+        """
+        关闭文件传输连接。
+
+        释放连接资源，应在不再需要连接时调用。
+        支持多次调用，子类应实现幂等性。
+        """
+        pass
+
+    @property
+    def name(self) -> str:
+        """
+        文件传输连接名称属性。
+
+        返回连接的名称，用于标识连接。
+
+        Returns:
+            str: 文件传输连接名称。
+        """
+        return self._name
+
+    @property
+    def logger(self) -> Logger:
+        """
+        日志记录器属性。
+
+        提供对内部日志记录器的只读访问，用于记录文件传输操作日志。
+
+        Returns:
+            Logger: Python logging 模块的日志记录器实例。
+        """
+        return self._logger
+
+    @property
+    def host(self) -> str:
+        """
+        主机地址属性。
+
+        Returns:
+            str: 文件服务器主机地址。
+        """
+        return self._host
+
+    @property
+    def port(self) -> int:
+        """
+        端口属性。
+
+        Returns:
+            int: 文件服务器端口号。
+        """
+        return self._port
+
+    @property
+    def user(self) -> str:
+        """
+        用户名属性。
+
+        Returns:
+            str: 文件服务器用户名。
+        """
+        return self._user
+
+    @property
+    def remote_path(self) -> str:
+        """
+        远程路径属性。
+
+        Returns:
+            str: 远程基础目录。
+        """
+        return self._remote_path
+
+    def __enter__(self):
+        """
+        上下文管理器入口方法。
+
+        使 FileTransferBase 的子类可以作为 with 语句使用。
+        进入 with 块时返回 self 引用。
+
+        Returns:
+            self: 返回当前实例的引用。
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        上下文管理器退出方法。
+
+        退出 with 块时自动调用，异常会被记录后关闭连接。
+        确保文件传输连接被正确关闭。
+
+        Args:
+            exc_type: 异常类型，若无异常则为 None。
+            exc_val: 异常值/实例，若无异常则为 None。
+            exc_tb: 异常回溯对象，若无异常则为 None。
+        """
+        if exc_type:
+            self._logger.exception(exc_val)
+        self.close()
+
+    def __del__(self):
+        """
+        析构函数。
+
+        当对象被垃圾回收时自动调用，确保连接被关闭。
+        这是一个安全网机制，不能替代显式调用 close()。
+        """
+        self.close()

dataops_sdk/runner/utils/ftp_file_transfer.py

@@ -0,0 +1,242 @@
+"""
+基于 FTP 的文件传输实现。
+
+继承自 FileTransferBase，使用 Python 内置的 ftplib 实现 FTP 文件传输。
+ftplib 是 Python 标准库，无需额外安装依赖。
+
+Classes:
+    FTPFileTransfer: 基于 ftplib 的 FTP 文件传输类
+
+Example:
+    >>> with FTPFileTransfer(host="ftp.example.com", port=21,
+    ...                     user="ftp_user", password="ftp_password",
+    ...                     remote_path="/uploads") as ftp:
+    ...     ftp.upload("/local/file.txt", "remote_file.txt")
+    ...     files = ftp.list_files("*.txt")
+    ...     if ftp.exists("remote_file.txt"):
+    ...         ftp.download("remote_file.txt", "/local/downloaded.txt")
+"""
+
+import fnmatch
+from ftplib import FTP, error_perm
+from io import BytesIO
+from pathlib import Path
+from typing import BinaryIO
+
+from .file_transfer_base import FileTransferBase
+
+
+class FTPFileTransfer(FileTransferBase):
+    """
+    基于 ftplib 的 FTP 文件传输类。
+
+    支持：
+    - 上传/下载文件
+    - 列出远程文件（支持 glob 模式过滤）
+    - 删除远程文件
+    - 检查文件是否存在
+    - 连接健康检查
+    - 上下文管理器（自动关闭连接）
+
+    继承自 FileTransferBase，实现所有抽象方法。
+
+    Attributes:
+        _ftp: ftplib.FTP 实例
+
+    Example:
+        >>> with FTPFileTransfer(host="ftp.example.com", port=21,
+        ...                     user="ftp_user", password="ftp_password",
+        ...                     remote_path="/uploads") as ftp:
+        ...     ftp.upload("/local/file.txt", "remote_file.txt")
+        ...     files = ftp.list_files("*.txt")
+    """
+
+    def __init__(
+        self,
+        host: str,
+        user: str,
+        password: str,
+        remote_path: str = "/",
+        port: int | str = 21,
+        name: str | None = None,
+        passive: bool = True,
+    ):
+        """
+        初始化 FTP 文件传输。
+
+        Args:
+            host: FTP 服务器主机地址。
+            user: FTP 服务器用户名。
+            password: FTP 服务器密码。
+            remote_path: 远程基础目录，默认为 "/"。
+            port: FTP 服务器端口号，默认为 21。
+            name: 连接名称，用于日志标识。
+            passive: 是否使用被动模式，默认为 True。
+                      True = 被动模式（客户端连接数据端口），适合服务器在防火墙后。
+                      False = 主动模式（服务器连接客户端数据端口），适合客户端在防火墙后。
+        """
+        super().__init__(host, user, password, remote_path, port, name)
+        self._passive = passive
+
+    def post_init(self):
+        """
+        初始化 FTP 连接。
+
+        连接到 FTP 服务器并切换到远程基础目录。
+
+        Returns:
+            None
+        """
+        self.logger.info(f"Connecting to FTP server: {self._host}:{self._port}")
+
+        self._ftp: FTP = FTP()
+        self._ftp.connect(self._host, self._port)
+        self._ftp.login(self._user, self._password)
+        self._ftp.keepalive()
+        self._ftp.set_pasv(self._passive)
+
+        # 切换到远程基础目录
+        self._ftp.cwd(self._remote_path)
+
+        self.logger.info(f"FTP connection established to {self._host}")
+
+    def upload(self, fp: BinaryIO | str | bytes, remote_name: str) -> bool:
+        """
+        上传文件到 FTP 服务器。
+
+        Args:
+            fp: 文件对象，支持任意 IO 类型（如 BytesIO, FileIO 等）。
+            remote_name: 远程文件名。
+
+        Returns:
+            bool: 上传成功返回 True，否则返回 False。
+        """
+        try:
+            if isinstance(fp, str):
+                fp = BytesIO(fp.encode("utf-8"))
+            elif isinstance(fp, bytes):
+                fp = BytesIO(fp)
+            self._ftp.storbinary(f"STOR {remote_name}", fp)
+            self.logger.info(f"Uploaded: -> {remote_name}")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to upload: {e}")
+            return False
+
+    def download(self, remote_name: str, local_path: str) -> bool:
+        """
+        从 FTP 服务器下载文件。
+
+        Args:
+            remote_name: 远程文件名。
+            local_path: 本地保存路径。
+
+        Returns:
+            bool: 下载成功返回 True，否则返回 False。
+        """
+        try:
+            # 确保本地目录存在
+            Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+
+            with open(local_path, "wb") as f:
+                self._ftp.retrbinary(f"RETR {remote_name}", f.write)
+            self.logger.info(f"Downloaded: {remote_name} -> {local_path}")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to download {remote_name}: {e}")
+            return False
+
+    def list_files(self, pattern: str | None = None) -> list[str]:
+        """
+        列出 FTP 远程目录下的文件。
+
+        Args:
+            pattern: 可选的 glob 模式，用于过滤文件名。
+
+        Returns:
+            list[str]: 文件名列表。
+        """
+        try:
+            files = self._ftp.nlst()
+            if pattern:
+                files = [f for f in files if fnmatch.fnmatch(f, pattern)]
+            return files
+        except error_perm as e:
+            # No files or empty directory
+            self.logger.debug(f"No files found or empty directory: {e}")
+            return []
+        except Exception as e:
+            self.logger.error(f"Failed to list files: {e}")
+            return []
+
+    def delete(self, remote_name: str) -> bool:
+        """
+        删除 FTP 服务器上的文件。
+
+        Args:
+            remote_name: 远程文件名。
+
+        Returns:
+            bool: 删除成功返回 True，否则返回 False。
+        """
+        try:
+            self._ftp.delete(remote_name)
+            self.logger.info(f"Deleted: {remote_name}")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to delete {remote_name}: {e}")
+            return False
+
+    def exists(self, remote_name: str) -> bool:
+        """
+        检查 FTP 服务器上文件是否存在。
+
+        Args:
+            remote_name: 远程文件名。
+
+        Returns:
+            bool: 文件存在返回 True，否则返回 False。
+        """
+        try:
+            # Use SIZE command - returns size if exists, raises error if not
+            self._ftp.size(remote_name)
+            return True
+        except error_perm:
+            return False
+        except Exception as e:
+            self.logger.error(f"Failed to check existence of {remote_name}: {e}")
+            return False
+
+    def ping(self) -> bool:
+        """
+        检查 FTP 连接是否存活。
+
+        Returns:
+            bool: 连接正常返回 True，否则返回 False。
+        """
+        try:
+            self._ftp.voidcmd("NOOP")
+            return True
+        except Exception as e:
+            self.logger.warning(f"FTP ping failed: {e}")
+            return False
+
+    def close(self):
+        """
+        关闭 FTP 连接。
+
+        关闭 FTP 连接，释放资源。
+        支持多次调用，子类应实现幂等性。
+        """
+        if hasattr(self, "_ftp") and self._ftp is not None:
+            try:
+                self._ftp.quit()
+                self.logger.info("FTP connection closed")
+            except Exception as e:
+                self.logger.error(f"Failed to close FTP connection: {e}")
+            self._ftp = None
+
+    @property
+    def passive(self) -> bool:
+        """被动模式"""
+        return self._passive

dataops_sdk/runner/utils/log.py

@@ -0,0 +1,50 @@
+"""日志模块。
+
+提供日志记录器创建功能，支持控制台和文件输出。
+文件日志按天轮转，保留 30 天。
+
+Functions:
+    getLogger: 获取或创建日志记录器
+    safe_str: 安全字符串转换，处理编码问题
+
+Example:
+    >>> logger = getLogger(__name__)
+    >>> logger.info("Hello, world!")
+"""
+
+import os
+from logging import INFO, Formatter, StreamHandler
+from logging import getLogger as _getLogger
+from logging.handlers import RotatingFileHandler
+
+from dataops_sdk.utils import get_log_path
+
+
+def getLogger(name: str | None = None, level: int = INFO):
+    _logger = _getLogger(name)
+    if _logger.hasHandlers():
+        return _logger
+
+    formatter = Formatter(
+        "%(asctime)s [%(levelname)s] %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    task_id = os.getenv("TASK_ID", None)
+    if task_id:
+        log_file = get_log_path(task_id)
+        # 文件处理器，按天分割（每天午夜轮转），保留30天的备份（约一个月）
+        # 使用 UTF-8 编码，避免 Unicode 编码错误
+        file_handler = RotatingFileHandler(
+            log_file, maxBytes=100 * 1024 * 1024, backupCount=1000, encoding="utf-8"
+        )
+        file_handler.setFormatter(formatter)
+        file_handler.setLevel(level)
+        _logger.addHandler(file_handler)
+
+    _handler = StreamHandler()
+    _handler.setFormatter(formatter)
+    _handler.setLevel(level)
+
+    _logger.addHandler(_handler)
+    _logger.setLevel(level)
+    return _logger

dataops_sdk/runner/utils/nacos.py

@@ -0,0 +1,96 @@
+"""Nacos 服务发现模块。
+
+提供从 Nacos 发现服务实例并返回 HTTP 客户端的功能。
+
+Functions:
+    get_service_instance: 从 Nacos 发现服务实例并返回 HTTP 客户端
+
+Example:
+    >>> client = await get_service_instance("my-service")
+    >>> response = await client.get("/api")
+"""
+
+import atexit
+import random
+from collections.abc import Mapping
+
+import httpx2
+from dataops_sdk.runner.schemas.nacos import NacosSettings, Service
+from dataops_sdk.runner.schemas.nacos.constants import Constants
+
+
+def get_service_instance(
+    service_name: str,
+    group_name: str = Constants.DEFAULT_GROUP,
+    clusters: list[str] | None = None,
+    headers: Mapping[str, str] | None = None,
+    nacos_addrs: str | None = None,
+    nacos_user: str | None = None,
+    nacos_password: str | None = None,
+    namespace_id: str | None = None,
+) -> httpx2.Client:
+    """
+    从Nacos发现服务实例并返回一个HTTP客户端
+
+    Args:
+        service_name (str): Nacos服务名称
+        group_name (str, None): Nacos服务分组，默认为 "DEFAULT_GROUP
+        clusters (list[str], None): Nacos服务集群，默认为空
+        headers: 可选的 HTTP 请求头
+        nacos_addrs: Nacos 服务地址，覆盖环境变量
+        nacos_user: Nacos 用户名，覆盖环境变量
+        nacos_password: Nacos 密码，覆盖环境变量
+        namespace_id: Nacos 命名空间 ID，覆盖环境变量
+
+    Returns:
+        httpx2.Client: 连接到选定服务实例的HTTP客户端
+
+    Raises:
+        NacosException: Nacos客户端错误(可能是连接问题、认证失败等)
+        RuntimeError: 如果没有可用的服务实例
+
+    Example:
+        >>> client = await get_service_instance("my-service")
+        >>> response = await client.get("/api")
+    """
+    if clusters is None:
+        clusters = []
+    config = NacosSettings()
+
+    if not nacos_user:
+        nacos_user = config.nacos_user
+
+    if not nacos_password:
+        nacos_password = config.nacos_password
+
+    if not nacos_addrs:
+        nacos_addrs = config.nacos_addrs
+
+    if not group_name:
+        group_name = config.nacos_group_name or Constants.DEFAULT_GROUP
+
+    if not namespace_id:
+        namespace_id = config.nacos_discovery_namespace
+    assert namespace_id, "namespace_id must be provided"
+
+    resp = httpx2.get(
+        f"http://{nacos_addrs}/nacos/v1/ns/instance/list",
+        params={
+            "username": nacos_user,
+            "password": nacos_password,
+            "serviceName": service_name,
+            "namespaceId": namespace_id,
+            "groupName": group_name,
+            "healthyOnly": "true",
+        },
+    )
+    service = Service.model_validate_json(resp.content)
+
+    instance = random.choice(service.hosts)
+
+    http_client = httpx2.Client(base_url=f"http://{instance.ip}:{instance.port}/{service_name}", headers=headers)
+
+    # 主进程结束前执行 http_client.close()
+    atexit.register(http_client.close)
+
+    return http_client