dataops-sdk 0.3.0__py3-none-any.whl

dataops_sdk/runner/utils/dbConn_base.py

@@ -0,0 +1,256 @@
+"""
+数据库连接基类模块。
+
+提供统一的数据库连接接口，支持 MySQL、PostgreSQL、DM、Oracle 等多种数据库。
+使用 SQLAlchemy 作为底层数据库抽象层，通过抽象基类定义统一的数据库操作规范。
+
+Example:
+    >>> class MySQLConn(DatabaseConnBase[MySQLConnection]):
+    ...     def post_init(self) -> MySQLConnection:
+    ...         engine = create_engine(self._url, echo=self._log_sql)
+    ...         return engine.connect()
+    ...
+    >>> with MySQLConn(task_id="task1", host="localhost", port=3306,
+    ...                user="root", password="pass", db="test") as conn:
+    ...     result = conn.select("SELECT * FROM users WHERE id = %s", (1,))
+    ...     conn.execute("INSERT INTO users (name) VALUES (%s)", ("Alice",))
+"""
+
+from abc import ABC, abstractmethod
+from logging import Logger
+from typing import Any
+
+from .log import getLogger
+
+
+class DatabaseConnBase(ABC):
+    """
+    数据库连接抽象基类。
+
+    定义数据库连接的标准接口，包括执行 SQL、批量操作、查询和关闭连接等方法。
+    使用泛型参数 C 表示具体数据库连接对象的类型（如 MySQLConnection、psycopg 连接等）。
+
+    继承此类需要实现以下抽象方法：
+    - post_init: 初始化具体连接对象
+    - executemany: 执行批量 SQL
+    - execute: 执行单条 SQL
+    - select: 查询数据
+    - close: 关闭连接
+
+    可以作为上下文管理器使用，会自动在退出时关闭连接。
+
+    Type Parameters:
+        C: 具体数据库连接对象的类型。
+
+    Attributes:
+        logger: 日志记录器，用于记录数据库操作日志。
+
+    Example:
+        >>> class MyDBConn(DatabaseConnBase[Any]):
+        ...     def post_init(self) -> Any:
+        ...         return object()  # 返回实际连接对象
+        ...     # ... 实现其他抽象方法
+        ...
+        >>> with MyDBConn(host="localhost", port=5432,
+        ...               user="user", password="pass", db="test") as conn:
+        ...     conn.execute("CREATE TABLE test (id INT)")
+        ...     conn.executemany("INSERT INTO test VALUES (%s)", [(1,), (2,)])
+        ...     rows = conn.select("SELECT * FROM test")
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int | str,
+        user: str,
+        password: str,
+        db: str,
+        db_type: str | None = "mysql",
+        log_sql: bool | None = None,
+        name: str | None = None,
+    ):
+        """
+        初始化数据库连接基类。
+
+        将传入的连接参数保存为实例属性，并初始化日志记录器。
+        子类应在 post_init 方法中创建实际的数据库连接。
+
+        Args:
+            task_id: 任务 ID，用于日志记录标识。
+            host: 数据库服务器主机地址。
+            port: 数据库服务器端口号，会被转换为整数。
+            user: 数据库用户名。
+            password: 数据库密码。
+            db: 数据库名称。
+            db_type: 数据库类型，默认为 "mysql"。
+            log_sql: 是否记录 SQL 语句日志，为 None 时使用实例级默认行为。
+        """
+        # 确保端口始终为整数类型
+        port = int(port)
+
+        # 保存数据库连接参数
+        self._host = host
+        self._port = int(port)  # 再次转换确保类型一致
+        self._user = user
+        self._password = password
+        self._db = db
+        self._db_type = db_type or "mysql"
+
+        self._name = name or f"{self._db_type}_{self._db}"
+
+        # log_sql 默认为 False，bool(None) 为 False，故用条件表达式处理
+        self._log_sql: bool = bool(log_sql) if log_sql is not None else False
+
+        # 使用 task_id 作为 logger 名称，便于日志追踪
+        self._logger = getLogger(self._name)
+
+        # 调用子类的初始化方法（钩子模式）
+        self.post_init()
+
+    @abstractmethod
+    def post_init(self):
+        """
+        初始化具体数据库连接对象。
+
+        子类应在此方法中创建并返回实际的数据库连接对象。
+        该方法在 __init__ 中被调用，用于延迟初始化连接。
+
+        Returns:
+            C: 具体类型的数据库连接对象。
+        """
+        pass
+
+    @abstractmethod
+    def executemany(
+        self,
+        sql: str,
+        datas: list[tuple | list],
+        *,
+        log_sql: bool | None = None,
+    ):
+        """
+        执行多条 SQL 语句（批量执行）。
+
+        用于批量插入、更新等场景，提高执行效率。
+
+        Args:
+            sql: 待执行的 SQL 语句模板，支持占位符（如 %s、? 等）。
+            datas: SQL 语句参数的列表，每一项对应一条 SQL 执行。
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql。
+        """
+        pass
+
+    @abstractmethod
+    def execute(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        *,
+        log_sql: bool | None = None,
+    ):
+        """
+        执行单条 SQL 语句。
+
+        可执行 INSERT、UPDATE、DELETE、CREATE 等语句。
+        默认不自动提交，需手动管理事务或配置自动提交。
+
+        Args:
+            sql: 待执行的 SQL 语句（可包含多条语句，用分号分隔）。
+            args: SQL 语句参数，支持元组、列表或字典格式。
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql。
+        """
+
+    @abstractmethod
+    def select(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        size: int | None = None,
+        *,
+        log_sql: bool | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        执行查询语句并返回结果。
+
+        用于 SELECT 查询场景，返回查询到的数据行。
+
+        Args:
+            sql: 待执行的查询 SQL 语句。
+            args: SQL 语句参数，支持元组、列表或字典格式。
+            size: 限制返回的行数，None 表示不限制。
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql。
+        """
+        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
+
+    def __enter__(self):
+        """
+        上下文管理器入口方法。
+
+        使 DatabaseConnBase 的子类可以作为 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()