dataops-sdk 0.3.0__py3-none-any.whl

dataops_sdk/runner/utils/neo4j_sync.py

@@ -0,0 +1,236 @@
+"""Neo4j 图数据库连接模块。
+
+提供 Neo4j 数据库连接管理，支持上下文管理器和 Cypher 语句执行。
+
+Classes:
+    Neo4jConn: Neo4j 图数据库连接类
+
+Example:
+    >>> with Neo4jConn(task_id="t1", host="localhost", port=7687,
+    ...               user="neo4j", password="123", db="neo4j") as conn:
+    ...     result = conn.select("MATCH (n) RETURN n LIMIT 10")
+"""
+
+from collections.abc import Iterable
+
+from .log import getLogger
+
+
+class Neo4jConn:
+    """Neo4j 图数据库连接类。
+
+    支持上下文管理器（with 语句）、Cypher 语句执行和结果查询。
+
+    Attributes:
+        driver: Neo4j 数据库驱动
+        session: Neo4j 数据库会话
+
+    Example:
+        >>> with Neo4jConn(task_id="t1", host="localhost", port=7687,
+        ...               user="neo4j", password="123", db="neo4j") as conn:
+        ...     result = conn.select("MATCH (n) RETURN n LIMIT 10")
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int | str,
+        user: str,
+        password: str,
+        db: str,
+        log_sql: bool | None = None,
+        name: str | None = None,
+    ):
+        from neo4j import GraphDatabase
+
+        self._host = host
+        self._port = port
+        self._user = user
+        self._password = password
+        self._database = db
+        # 实例级 log_sql 默认值：None 表示关闭（与 LoggerConfig.sql=False 一致）；
+        # 各方法的 log_sql=None 时使用此值；显式传 bool 则覆盖。
+        self._log_sql: bool = bool(log_sql) if log_sql is not None else False
+        self._driver = GraphDatabase.driver(
+            f"bolt://{host}:{port}",
+            auth=(user, password),
+            database=db,
+        )
+        self._session = self._driver.session()
+        self._name = name or f"neo4j_{host}"
+        self._logger = getLogger(self._name)
+
+    def _should_log_sql(self, log_sql: bool | None) -> bool:
+        """
+        解析方法级 log_sql 参数：
+        - None（未传）→ 使用实例级 self._log_sql
+        - 显式 bool → 使用调用方值
+        """
+        return self._log_sql if log_sql is None else bool(log_sql)
+
+    def executemany(
+        self,
+        sql: str,
+        datas: Iterable[tuple | list],
+        *,
+        auto_commit: bool = False,
+        log_sql: bool | None = None,
+        **kwargs,
+    ):
+        """
+        执行多条 Cypher 语句
+
+        Args:
+            sql: 待执行的 Cypher 语句
+            datas: 待执行的语句参数列表
+            auto_commit: 是否自动提交，默认为 True
+            log_sql: 是否打印SQL日志，默认为 None（使用实例级 self._log_sql）
+        """
+        if self._should_log_sql(log_sql):
+            self._logger.debug(f"executemany:\n {sql}")
+        for data in datas:
+            if hasattr(self, "_transaction") and self._transaction is not None:
+                self._transaction.run(sql, data)
+            else:
+                self._session.run(sql, data)
+
+    def execute(
+        self,
+        sql: str,
+        args: dict | list | tuple | None = None,
+        *,
+        auto_commit: bool = False,
+        log_sql: bool | None = None,
+        **kwargs,
+    ):
+        """
+        执行 Cypher 语句
+
+        Args:
+            sql: 待执行的 Cypher 语句
+            args: 待执行的语句参数
+            auto_commit: 是否自动提交，默认为 True
+            log_sql: 是否打印SQL日志，默认为 None（使用实例级 self._log_sql）
+        """
+        if self._should_log_sql(log_sql):
+            self._logger.debug(f"execute:\n {sql}")
+        if hasattr(self, "_transaction") and self._transaction is not None:
+            self._transaction.run(sql, args)
+        else:
+            self._session.run(sql, args)
+
+    def select(
+        self,
+        sql: str,
+        args: dict | list | tuple | None = None,
+        size: int | None = None,
+        *,
+        log_sql: bool | None = None,
+        **kwargs,
+    ) -> list[dict]:
+        """
+        查询数据
+
+        Args:
+            cypher: 待执行的 Cypher 查询语句
+            args: 查询参数
+            size: 返回结果数量限制，None 表示返回全部
+            log_sql: 是否打印SQL日志，默认为 None（使用实例级 self._log_sql）
+
+        Returns:
+            查询结果列表
+        """
+        if self._should_log_sql(log_sql):
+            self._logger.debug(f"select:\n {sql}")
+        if hasattr(self, "_transaction") and self._transaction is not None:
+            result = self._transaction.run(sql, args).data()
+        else:
+            result = self._session.run(sql, args).data()
+
+        if size is not None and size > 0:
+            result = result[:size]
+
+        return result
+
+    # def begin(self):
+    #     """
+    #     开始事务
+    #     """
+    #     if not hasattr(self, "_transaction") or self._transaction is None:
+    #         self._transaction = self._session.begin_transaction()
+
+    # def commit(self):
+    #     """
+    #     提交事务
+    #     """
+    #     if hasattr(self, "_transaction") and self._transaction is not None:
+    #         self._transaction.commit()
+    #         self._transaction = None
+
+    # def rollback(self):
+    #     """
+    #     回滚事务
+    #     """
+    #     if hasattr(self, "_transaction") and self._transaction is not None:
+    #         self._transaction.rollback()
+    #         self._transaction = None
+
+    def close(self):
+        """
+        关闭数据库连接
+        """
+        if hasattr(self, "_session") and self._session is not None:
+            self._session.close()
+        if hasattr(self, "_driver") and self._driver is not None:
+            self._driver.close()
+
+    def __enter__(self):
+        """
+        上下文管理器入口
+
+        Returns:
+            self
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val: Exception, exc_tb):
+        """
+        上下文管理器退出，自动关闭连接
+
+        Args:
+            exc_type: 异常类型
+            exc_val: 异常值
+            exc_tb: 异常回溯
+        """
+        self.close()
+        if exc_type:
+            raise exc_val
+
+    def __del__(self):
+        """
+        析构函数，确保连接被关闭
+        """
+        self.close()
+
+    @property
+    def driver(self):
+        """
+        获取数据库连接
+
+        Returns:
+            数据库连接
+        """
+        return self._driver
+
+    def get_session(self):
+        """
+        获取数据库会话
+
+        Returns:
+            数据库会话
+        """
+        return self._driver.session()
+
+    @property
+    def name(self) -> str:
+        return self._name

dataops_sdk/runner/utils/queue_base.py

@@ -0,0 +1,224 @@
+"""
+队列操作抽象基类模块。
+
+提供统一的队列操作接口，支持 Kafka、RabbitMQ、Redis 等多种消息队列。
+使用抽象基类定义统一的队列操作规范。
+
+Example:
+    >>> class KafkaQueue(QueueBase):
+    ...     def post_init(self):
+    ...         from kafka import KafkaProducer, KafkaConsumer
+    ...         self._producer = KafkaProducer(bootstrap_servers=self._bootstrap_servers)
+    ...         self._consumer = KafkaConsumer(self._topic,
+    ...                                        bootstrap_servers=self._bootstrap_servers,
+    ...                                        group_id=self._group_id)
+    ...     def send(self, value, key=None):
+    ...         self._producer.send(self._topic, value=value, key=key)
+    ...         self._producer.flush()
+    ...     def recv(self, timeout=None):
+    ...         msg = self._consumer.poll(timeout_ms=int((timeout or 1) * 1000))
+    ...         for tp, msgs in msg.items():
+    ...             for m in msgs:
+    ...                 return m.value, m.key
+    ...         return None
+    ...     def close(self):
+    ...         self._producer.close()
+    ...         self._consumer.close()
+    ...
+    >>> with KafkaQueue(task_id="task1", bootstrap_servers="localhost:9092",
+    ...                 topic="test", group_id="test-group") as queue:
+    ...     queue.send(b"hello", key=b"k1")
+    ...     msg = queue.recv(timeout=5)
+"""
+
+from abc import ABC, abstractmethod
+from logging import Logger
+
+from .log import getLogger
+
+
+class QueueBase(ABC):
+    """
+    队列操作抽象基类。
+
+    定义队列操作的标准接口，包括发送消息、接收消息和关闭连接等方法。
+    使用抽象基类定义统一的队列操作规范。
+
+    继承此类需要实现以下抽象方法：
+    - post_init: 初始化具体队列连接对象
+    - send: 发送消息
+    - recv: 接收消息
+    - close: 关闭连接
+
+    可以作为上下文管理器使用，会自动在退出时关闭连接。
+
+    Attributes:
+        logger: 日志记录器，用于记录队列操作日志。
+
+    Example:
+        >>> class MyQueue(QueueBase):
+        ...     def post_init(self):
+        ...         pass  # 初始化具体连接
+        ...     def send(self, value, key=None):
+        ...         pass
+        ...     def recv(self, timeout=None):
+        ...         return None
+        ...     def close(self):
+        ...         pass
+        ...
+        >>> with MyQueue(task_id="t1", bootstrap_servers="localhost:9092",
+        ...              topic="test") as queue:
+        ...     queue.send(b"hello")
+        ...     msg = queue.recv(timeout=5)
+    """
+
+    def __init__(
+        self,
+        bootstrap_servers: str,
+        name: str | None = None,
+    ):
+        """
+        初始化队列基类。
+
+        将传入的连接参数保存为实例属性，并初始化日志记录器。
+        子类应在 post_init 方法中创建实际的队列连接。
+
+        Args:
+            bootstrap_servers: 消息队列服务器地址。
+            name: 名称。
+        """
+        self._bootstrap_servers = bootstrap_servers
+
+        self._name = name or "queue_unknown"
+
+        self._logger = getLogger(name)
+
+        self.post_init()
+
+    @abstractmethod
+    def post_init(self):
+        """
+        初始化具体队列连接对象。
+
+        子类应在此方法中创建并返回实际的队列连接对象。
+        该方法在 __init__ 中被调用，用于延迟初始化连接。
+
+        Returns:
+            None
+        """
+        pass
+
+    @abstractmethod
+    def send(self, topic: str, value: bytes | str, key: bytes | str | None = None):
+        """
+        发送消息到队列。
+
+        将消息发送到指定的主题。
+
+        Args:
+            topic: 消息主题。
+            value: 消息内容，支持字节串或字符串。
+            key: 消息键，用于分区路由，可选。
+        """
+        pass
+
+    @abstractmethod
+    def recv(
+        self, topic: str, group_id: str | None = None, timeout: float | None = None
+    ) -> tuple[bytes | str, bytes | str] | None:
+        """
+        从队列接收消息。
+
+        从指定主题消费消息。
+
+        Args:
+            topic: 消息主题。
+            group_id: 消费者组 ID，可选；为 None 时使用初始化时的 group_id。
+            timeout: 等待超时时间（秒），None 表示阻塞等待。
+
+        Returns:
+            tuple[bytes | str, bytes | str] | None: 消息的 (value, key) 元组，无消息时返回 None。
+        """
+        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 bootstrap_servers(self) -> str:
+        """
+        服务器地址属性。
+
+        Returns:
+            str: 队列服务器地址。
+        """
+        return self._bootstrap_servers
+
+
+    def __enter__(self):
+        """
+        上下文管理器入口方法。
+
+        使 QueueBase 的子类可以作为 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/queue_kafka.py

@@ -0,0 +1,150 @@
+"""
+基于 Kafka 的消息队列实现。
+
+继承自 QueueBase，提供统一的队列操作接口。
+使用 kafka-python 作为底层实现，支持消息的可靠发送和消费者组模式。
+
+Classes:
+    KafkaQueue: 基于 kafka-python 的 Kafka 消息队列类
+
+Example:
+    >>> with KafkaQueue(task_id="t1", bootstrap_servers="localhost:9092",
+    ...                 topic="test", group_id="test-group") as queue:
+    ...     queue.send(b"hello", key=b"k1")
+    ...     msg = queue.recv(timeout=5)
+    ...     if msg:
+    ...         value, key = msg
+    ...         print(f"Received: {value}")
+"""
+
+from kafka import KafkaConsumer, KafkaProducer
+from kafka.errors import KafkaError
+
+from .queue_base import QueueBase
+
+__all__ = ["KafkaQueue"]
+
+
+class KafkaQueue(QueueBase):
+    """
+    基于 kafka-python 的 Kafka 消息队列类。
+
+    支持：
+    - 可靠发送消息（等待 broker 确认）
+    - 消费者组模式（多实例自动负载均衡）
+    - 上下文管理器（自动关闭连接）
+
+    继承自 QueueBase，实现所有抽象方法。
+
+    Attributes:
+        producer: KafkaProducer 实例。
+        logger: 日志记录器。
+
+    Example:
+        >>> with KafkaQueue(task_id="t1", bootstrap_servers="localhost:9092",
+        ...                 group_id="test-group") as queue:
+        ...     # 发送消息
+        ...     queue.send("test-topic", b"hello world", key=b"greeting")
+        ...     # 接收消息
+        ...     msg = queue.recv(timeout=5.0)
+        ...     if msg:
+        ...         value, key = msg
+        ...         print(f"key={key}, value={value}")
+    """
+
+    def post_init(self):
+        """
+        初始化 KafkaProducer。
+
+        Returns:
+            None
+        """
+        self.logger.info(f"Creating Kafka producer: bootstrap_servers={self._bootstrap_servers}")
+
+        self._producer: KafkaProducer = KafkaProducer(
+            bootstrap_servers=self._bootstrap_servers.split(","),
+            acks=1,
+            retries=3,
+            max_in_flight_requests_per_connection=1,
+        )
+
+    def send(self, topic: str, value: bytes | str, key: bytes | str | None = None):
+        """
+        发送消息到 Kafka 主题。
+
+        使用同步发送模式，等待 broker 确认消息已写入。
+
+        Args:
+            topic: 消息主题。
+            value: 消息内容，支持字节串或字符串。
+            key: 消息键，用于分区路由，可选。
+        """
+        try:
+            if isinstance(value, str):
+                value = value.encode("utf-8")
+            if isinstance(key, str):
+                key = key.encode("utf-8")
+
+            future = self._producer.send(topic, value=value, key=key)
+            record_metadata = future.get(timeout=10)
+
+            self.logger.info(
+                f"Message sent to {record_metadata.topic}:{record_metadata.partition}:"
+                f"{record_metadata.offset}"
+            )
+        except KafkaError as e:
+            self.logger.error(f"Failed to send message: {e}")
+            raise
+
+    def recv(
+        self, topic: str, group_id: str | None = None, timeout: float | None = None
+    ) -> tuple[bytes, bytes] | None:
+        """
+        从 Kafka 主题接收消息。
+
+        使用消费者组模式，支持多实例自动负载均衡。
+
+        Args:
+            topic: 消息主题。
+            group_id: 消费者组 ID，可选；为 None 时使用初始化时的 group_id。
+            timeout: 等待超时时间（秒），None 表示使用默认超时（1秒）。
+
+        Returns:
+            tuple[bytes, bytes] | None: 消息的 (value, key) 元组，无消息时返回 None。
+        """
+        try:
+            with KafkaConsumer(
+                topic,
+                bootstrap_servers=self._bootstrap_servers.split(","),
+                group_id=group_id,
+                auto_offset_reset="earliest",
+                enable_auto_commit=True,
+            ) as consumer:
+                msg = consumer.poll(timeout_ms=int((timeout or 1.0) * 1000))
+
+                for _, messages in msg.items():
+                    for m in messages:
+                        return m.value, m.key
+
+                return None
+        except Exception as e:
+            self.logger.error(f"Failed to receive message: {e}")
+            return None
+
+    def close(self):
+        """
+        关闭 Kafka 连接。
+
+        关闭 Producer，释放资源。
+        支持多次调用，子类应实现幂等性。
+        """
+        if hasattr(self, "_producer") and self._producer is not None:
+            try:
+                self._producer.flush()
+                self._producer.close()
+                self.logger.info("Producer closed")
+            except Exception as e:
+                self.logger.error(f"Failed to close producer: {e}")
+            self._producer = None
+
+        self.logger.info("closed")