dataops-sdk 0.3.0__py3-none-any.whl

dataops_sdk/runner/utils/dbConn_sync.py

@@ -0,0 +1,485 @@
+"""
+基于 SQLAlchemy 的同步数据库连接实现。
+
+继承自 DatabaseConnBase，提供统一的数据库连接接口。
+使用 SQLAlchemy Engine 和 Session 作为底层实现，支持 MySQL、PostgreSQL、DM、Oracle 等多种数据库。
+
+Classes:
+    DatabaseConn: 基于 SQLAlchemy 的同步数据库连接类
+
+Example:
+    >>> with DatabaseConn(task_id="t1", 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",))
+    ...     conn.executemany("INSERT INTO users (name) VALUES (%s)", [("Bob",), ("Carol",)])
+"""
+
+import re
+from collections.abc import Mapping
+from typing import Any, Final, Literal, overload
+
+from sqlalchemy import URL, Engine, create_engine, text
+from sqlalchemy.orm import Session, sessionmaker
+
+from .dbConn_base import DatabaseConnBase
+
+# 数据库类型到驱动程序的映射
+# 用于 build_url 函数构建数据库连接 URL
+DBTYPE_DRIVER: Final[Mapping[str, str]] = {
+    "mysql": "pymysql",
+    "postgresql": "psycopg",
+    "dm": "dmpython",
+    "oracle": "oracledb",
+    "doris": "pymysql",
+}
+
+
+def build_url(
+    host: str,
+    port: int | str,
+    user: str,
+    password: str,
+    db: str,
+    db_type: str,
+) -> str | URL:
+    """
+    构建数据库连接 URL 字符串。
+
+    将数据库连接参数格式化为 SQLAlchemy 所使用的 URL 格式。
+    用户名和密码会进行 URL 编码，以处理特殊字符。
+
+    Args:
+        host: 数据库服务器主机地址。
+        port: 数据库服务器端口号。
+        user: 数据库用户名。
+        password: 数据库密码。
+        db: 数据库名称。
+        db_type: 数据库类型
+
+    Returns:
+        格式化的数据库连接 URL 字符串。
+
+    Raises:
+        NotImplementedError: 当指定的 db_type 不在支持列表中时抛出。
+
+    Example:
+        >>> url = build_url("localhost", 3306, "root", "password", "mydb", "mysql")
+        >>> print(url)
+        mysql+pymysql://root:password@localhost:3306/mydb
+    """
+    # 检查数据库类型是否支持
+    if db_type not in DBTYPE_DRIVER:
+        raise NotImplementedError(f"{db_type} 尚未支持，请联系管理员")
+
+    query: dict[str, str] = dict()
+    if db_type in {"mysql", "doris"}:
+        query["local_infile"] = "True"
+        query["ssl_disabled"] = "True"
+
+    driver = DBTYPE_DRIVER[db_type]
+    drivername = f"{db_type}+{driver}" if driver else db_type
+    return URL.create(
+        drivername=drivername,
+        username=user,
+        password=password,
+        host=host,
+        port=int(port),
+        database=db,
+        query=query,
+    )
+
+
+class DatabaseConn(DatabaseConnBase):
+    """
+    基于 SQLAlchemy 的同步数据库连接类。
+
+    使用 SQLAlchemy Engine 和 Session 实现数据库操作，支持：
+    - 单条 SQL 执行（execute）
+    - 批量 SQL 执行（executemany）
+    - 查询操作（select）
+    - 上下文管理器（自动关闭连接）
+    - 事务管理（commit/rollback）
+
+    继承自 DatabaseConnBase，实现所有抽象方法。
+
+    Attributes:
+        engine: SQLAlchemy Engine 对象。
+        session: SQLAlchemy Session 对象。
+        logger: 日志记录器。
+
+    Example:
+        >>> with DatabaseConn(task_id="t1", host="localhost", port=3306,
+        ...                       user="root", password="pass", db="test",
+        ...                       db_type="mysql") as conn:
+        ...     # 查询数据
+        ...     rows = conn.select("SELECT * FROM users LIMIT %s", (10,))
+        ...     # 插入数据
+        ...     conn.execute("INSERT INTO users (name, email) VALUES (%s, %s)",
+        ...                   ("Alice", "alice@example.com"))
+        ...     # 批量插入
+        ...     conn.executemany("INSERT INTO users (name) VALUES (%s)",
+        ...                       [("Bob",), ("Carol",)])
+    """
+
+    def post_init(self):
+        """
+        初始化 SQLAlchemy Engine 和 Session。
+
+        使用 build_url 构建数据库 URL，创建 Engine 和 SessionFactory。
+
+        Returns:
+            Session: SQLAlchemy Session 对象。
+        """
+        # 构建数据库连接 URL
+        self._url = build_url(
+            host=self._host,
+            port=self._port,
+            user=self._user,
+            password=self._password,
+            db=self._db,
+            db_type=self._db_type,
+        )
+
+        # 创建 SQLAlchemy Engine
+        if self._log_sql:
+            self._logger.info(
+                f"Creating SQLAlchemy Engine for {self._host}:{self._port}/{self._db}"
+            )
+        self._engine: Engine = create_engine_safe(self._db_type, self._url)
+
+        # 创建 SessionFactory
+        self._session_factory = sessionmaker(
+            bind=self._engine,
+            autoflush=False,
+            expire_on_commit=False,
+            autocommit=False,
+        )
+        self._session = None
+
+    def _get_log_sql(self, log_sql: bool | None) -> bool:
+        """
+        获取是否记录 SQL 的布尔值。
+
+        优先级：方法参数 log_sql > 实例级 self._log_sql
+
+        Args:
+            log_sql: 方法级 log_sql 参数。
+
+        Returns:
+            bool: 是否应记录 SQL。
+        """
+        return bool(log_sql) if log_sql is not None else self._log_sql
+
+    @staticmethod
+    def render_sql(sql: str, args: tuple | list | dict | None) -> str:
+        """
+        渲染 SQL 语句。
+
+        将 SQL 语句和参数组合成完整的 SQL 语句。
+
+        Args:
+            sql: SQL 语句模板。
+            args: SQL 语句参数。
+
+        Returns:
+            str: 渲染后的 SQL 语句。
+        """
+        if not args:
+            return sql
+        if isinstance(args, (list, tuple)):
+            return sql.replace("%s", "{}").format(*args)
+        else:
+            # 把 :name 替换为 {name}
+            sql = re.sub(r"\%\((\w+)\)s", r"\{\1\}", sql)
+            return sql.format(**args)
+
+    def executemany(
+        self,
+        sql: str,
+        datas: list[tuple | list],
+        *,
+        auto_commit: bool = True,
+        log_sql: bool | None = None,
+        **kwargs,
+    ):
+        """
+        执行多条 SQL 语句（批量执行）。
+
+        用于批量插入、更新等场景。会自动提交事务。
+
+        Args:
+            sql: 待执行的 SQL 语句模板，支持 %s 或 :name 占位符。
+            datas: SQL 语句参数的列表，每一项对应一条 SQL 执行。
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql。
+        """
+        if kwargs:
+            self.logger.warning(f"尚未支持的参数: {kwargs}")
+
+        should_log = self._get_log_sql(log_sql)
+
+        sql, args = self.process_input(sql, args)
+
+        if should_log:
+            self.logger.info(self.render_sql(sql, args[0]))
+
+        if self.session:
+            self.session.execute(text(sql), datas)
+
+            if auto_commit:
+                session.commit()
+        else:
+            with self.session_factory() as session:
+                session.execute(text(sql), datas)
+
+                if auto_commit:
+                    session.commit()
+
+    def execute(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        *,
+        auto_commit: bool = True,
+        log_sql: bool | None = None,
+        **kwargs,
+    ):
+        """
+        执行单条 SQL 语句。
+
+        可执行 INSERT、UPDATE、DELETE、CREATE 等语句。
+        默认自动提交事务。
+
+        Args:
+            sql: 待执行的 SQL 语句。
+            args: SQL 语句参数，支持元组、列表或字典格式。
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql。
+        """
+        if kwargs:
+            self.logger.warning(f"尚未支持的参数: {kwargs}")
+
+        should_log = self._get_log_sql(log_sql)
+
+        sql, args = self.process_input(sql, args)
+
+        if should_log:
+            self.logger.info(self.render_sql(sql, args))
+
+        if self.session:
+            self.session.execute(text(sql), args)
+
+            if auto_commit:
+                session.commit()
+        else:
+            with self.session_factory() as session:
+                session.execute(text(sql), args)
+
+            if auto_commit:
+                session.commit()
+
+    @overload
+    def select(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        size: int | None = None,
+        mapping: Literal[True] = True,
+        *,
+        log_sql: bool | None = None,
+    ) -> list[dict[str, Any]]:
+        ...
+
+    @overload
+    def select(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        size: int | None = None,
+        mapping: Literal[False] = False,
+        *,
+        log_sql: bool | None = None,
+    ) -> list[tuple[Any]]:
+        ...
+
+    def select(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        size: int | None = None,
+        mapping: bool = False,
+        *,
+        log_sql: bool | None = None,
+    ) -> list[dict[str, Any] | tuple[Any]]:
+        """
+        执行查询语句并返回结果。
+
+        用于 SELECT 查询场景，返回查询到的数据行列表。
+
+        Args:
+            sql: 待执行的查询 SQL 语句。
+            args: SQL 语句参数，支持元组、列表或字典格式。
+            size: 限制返回的行数，None 表示不限制。
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql。
+
+        Returns:
+            list[tuple]: 查询结果，元素为 Row 对象或字典（取决于 SQLAlchemy 配置）。
+        """
+        should_log = self._get_log_sql(log_sql)
+
+        if isinstance(sql, str):
+            sql, args = self.process_input(sql, args)
+
+        if should_log:
+            self.logger.info(self.render_sql(sql, args))
+
+        with self.session_factory() as session, session.connection() as conn:
+            result = (
+                conn.exec_driver_sql(sql, args)
+                if isinstance(sql, str)
+                else sql.execute(session.connection())
+            )
+
+            if mapping:
+                result = result.mappings()
+            # 获取结果
+            rows = result.fetchall() if size is None or size <= 0 else result.fetchmany(size)
+
+        return list(map(dict, rows)) if mapping else list(map(lambda x: x._t._data, rows))
+
+    @staticmethod
+    def process_input(sql: str, args: dict | tuple | list | None):
+        # sql = sql.replace("%s", "?")
+        # # 把 %(key)s 替换为 :key
+        # sql = re.sub(r"\%\((\w+)\)s", r":\1", sql)
+        if args is None:
+            return sql, None
+        if isinstance(args, list):
+            if len(args) > 1 and isinstance(args[0], list):
+                return sql, list(map(tuple, args))
+            else:
+                return sql, tuple(args)
+                # return sql, args
+        else:
+            return sql, args
+
+    def close(self):
+        """
+        关闭数据库连接。
+
+        关闭 Session 和 Engine，释放数据库连接资源。
+        支持多次调用，子类应实现幂等性。
+        """
+        # 关闭 Session
+        if hasattr(self, "_session") and self._session is not None:
+            self._session.close()
+            self._session = None
+
+        # 关闭 Engine（会关闭所有连接池）
+        if hasattr(self, "_engine") and self._engine is not None:
+            self.engine.dispose()
+            self._engine = None
+            self._logger.info(f"closed")
+
+    @property
+    def session_factory(self):
+        return self._session_factory
+
+    @property
+    def engine(self) -> Engine:
+        """
+        获取 SQLAlchemy Engine 对象。
+
+        Returns:
+            Engine: SQLAlchemy Engine 对象。
+        """
+        return self._engine
+
+    def commit(self):
+        """
+        提交当前事务。
+        """
+        self._logger.debug(f"committing")
+        self.session.commit()
+
+    def rollback(self):
+        """
+        回滚当前事务。
+        """
+        self._logger.warning(f"rolling back")
+        self.session.rollback()
+
+    def begin(self):
+        """
+        开始一个新的事务。
+        """
+        self._logger.debug(f"beginning")
+        if self.session:
+            assert self._session is not None, "session not closed"
+        self._session = self.session_factory()
+        self._session.begin()
+
+    @property
+    def session(self) -> Session | None:
+        """
+        获取当前事务的 Session 对象。
+
+        Returns:
+            Session: 当前事务的 Session 对象。
+        """
+        return self._session
+
+class DatabaseConnDataOps(DatabaseConn):
+    def __init__(self, datasource: str, log_sql: bool | None = None):
+        from dataops_sdk.services.meta import Meta, release_list_get_by_name
+
+        datasource_meta = Meta.model_validate(release_list_get_by_name("datasource", datasource))
+
+        super().__init__(
+            task_id="",
+            host=datasource_meta.data["ip"],
+            port=datasource_meta.data["port"],
+            user=datasource_meta.data["username"],
+            password=datasource_meta.data["password"],
+            db=datasource_meta.data["database"],
+            db_type=datasource_meta.data["type"],
+            log_sql=log_sql,
+        )
+
+
+def create_engine_safe(db_type: str, url: str, **kwargs: Any) -> Engine:
+    """
+    安全创建 SQLAlchemy Engine。
+
+    处理各种数据库驱动的兼容性问题。
+
+    Args:
+        url: 数据库连接 URL。
+        **kwargs: 传递给 create_engine 的参数。
+
+    Returns:
+        Engine: SQLAlchemy Engine 对象。
+
+    Raises:
+        NotImplementedError: 当数据库类型不支持时抛出。
+    """
+    global DBTYPE_DRIVER
+
+    if db_type not in DBTYPE_DRIVER:
+        raise NotImplementedError(f"{db_type} 尚未支持，请联系管理员")
+
+    # 根据不同数据库类型调整连接参数
+    if db_type in ("mysql", "doris"):
+        # MySQL 需要添加一些兼容参数
+        kwargs.setdefault("pool_pre_ping", True)
+        kwargs.setdefault("pool_recycle", 3600)
+    elif db_type == "postgresql":
+        # PostgreSQL 配置
+        kwargs.setdefault("pool_pre_ping", True)
+    elif db_type == "dm":
+        # DM 数据库配置
+        kwargs.setdefault("pool_pre_ping", True)
+    elif db_type == "oracle":
+        # Oracle 配置
+        kwargs.setdefault("pool_pre_ping", True)
+
+    return create_engine(url, **kwargs)

dataops_sdk/runner/utils/duckdb_manager.py

@@ -0,0 +1,203 @@
+"""DuckDB 管理模块。
+
+提供 DuckDB 连接和生命周期管理，支持从传统数据库导入数据到 DuckDB。
+
+Classes:
+    DuckDBManager: DuckDB 连接和生命周期管理器
+
+Example:
+    >>> with DuckDBManager('task1') as conn:
+    ...     conn.sql("SELECT * FROM table").fetchall()
+    # 自动关闭连接并删除临时文件
+"""
+
+import asyncio
+from collections.abc import Coroutine
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+import duckdb
+import pandas as pd
+
+from .dbConn_sync import DatabaseConn
+from .log import getLogger
+
+
+class DuckDBManager:
+    """
+    DuckDB 连接和生命周期管理器
+
+    支持上下文管理器（with语法），自动管理连接和临时文件清理
+
+    使用示例：
+        with DuckDBManager('port') as conn:
+            conn.sql("SELECT * FROM table").fetchall()
+        # 自动关闭连接并删除临时文件
+    """
+
+    def __init__(self, task_id: str, prefix: str = "duckdb"):
+        """
+        初始化 DuckDB 管理器
+
+        Args:
+            prefix: 临时文件前缀，默认 'duckdb'
+        """
+        self._conn: duckdb.DuckDBPyConnection | None = None
+        self._temp_file = Path(f"{prefix}_{uuid4().hex}.db")
+        self._logger = getLogger(task_id)
+        self._loop = None
+
+    def _run_coroutine(self, coro: Coroutine):
+        """同步执行异步协程"""
+        if self._loop is None:
+            loop = asyncio.get_event_loop()
+            if loop is not None:
+                self._loop = loop
+            else:
+                self._loop = asyncio.new_event_loop()
+            assert self._loop, "no running loop"
+
+        return self._loop.run_until_complete(coro)
+
+    def __enter__(self) -> duckdb.DuckDBPyConnection:
+        """进入上下文，创建连接"""
+        self._conn = duckdb.connect(self._temp_file)
+        self._logger.info(f"DuckDB connection created: {self._temp_file}")
+        return self._conn
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """退出上下文，关闭连接并清理临时文件"""
+        self.close()
+        return False
+
+    @property
+    def conn(self) -> duckdb.DuckDBPyConnection | None:
+        """获取连接对象"""
+        return self._conn
+
+    @property
+    def is_connected(self) -> bool:
+        """检查是否已连接"""
+        return self._conn is not None
+
+    def execute(self, sql: str, parameters: object = None) -> duckdb.DuckDBPyConnection:
+        """执行 SQL 语句"""
+        self.open()
+        return self._conn.execute(sql, parameters=parameters)
+
+    def sql(
+        self, sql: duckdb.Statement | str, *, alias: str = "", params: object = None
+    ) -> duckdb.DuckDBPyConnection:
+        """执行 SQL 查询并返回结果"""
+        self.open()
+        return self._conn.sql(sql, alias=alias, params=params)
+
+    def close(self):
+        """关闭连接并清理临时文件"""
+        if self._conn:
+            try:
+                self._conn.close()
+                self._logger.info("DuckDB connection closed")
+            except Exception as e:
+                self._logger.warning(f"Error closing DuckDB connection: {e}")
+            self._conn = None
+
+        if self._temp_file and self._temp_file.exists():
+            try:
+                self._temp_file.unlink()
+                self._logger.info(f"Temporary file {self._temp_file} deleted")
+            except Exception as e:
+                self._logger.warning(f"Error deleting temporary file {self._temp_file}: {e}")
+
+    def open(self) -> duckdb.DuckDBPyConnection:
+        """打开连接（推荐使用 with 语句）"""
+        if not self._conn:
+            self._conn = duckdb.connect(self._temp_file)
+            self._logger.info(f"DuckDB connection created: {self._temp_file}")
+        return self._conn
+
+    def import_from_db(
+        self,
+        db_conn: DatabaseConn,
+        sql: str,
+        table_name: str,
+        page_size: int = 200000,
+    ):
+        """
+        从数据库导入数据到 duckdb
+
+        Args:
+            duck_db_conn: duckdb 数据库连接
+            db_conn: 数据库连接
+            sql: SQL 语句
+            table_name: duckdb 表名
+            page_size: 分页大小
+        """
+        self.open()
+        page_num = 1
+        total_count = 0
+        db_conn.execute(sql)
+
+        # 从cursor.description获取列名
+        columns = [col[0] for col in db_conn.cursor.description]
+
+        if db_conn.cursor.rowcount == 0:
+            # 如果没有数据，需要在 duckdb 创建表结构
+            db_conn.logger.info("No data from query, creating table structure in DuckDB")
+            field_str = ", ".join([f"{col} VARCHAR" for col in columns])
+            sql_create_table = f"CREATE TABLE IF NOT EXISTS {table_name} ({field_str})"
+            db_conn.logger.info(f"Create table SQL: {sql_create_table}")
+            # await asyncio.to_thread(duck_db_conn.execute, sql_create_table)
+            self.execute(sql_create_table)
+        else:
+            while True:
+                db_conn.logger.info(f"Processing page {page_num}")
+                result_list = db_conn.cursor.fetchmany(page_size)
+                df = pd.DataFrame(result_list, columns=columns)
+                if len(df) == 0:
+                    break
+                if page_num == 1:
+                    # 第一次导入数据时，先创建表结构
+                    # 将 DataFrame 注册为临时视图，以便创建表结构
+                    self._conn.register("df_view", df)
+                    self.execute(
+                        f"CREATE TABLE IF NOT EXISTS {table_name} AS SELECT * FROM df_view WHERE 1=0"
+                    )
+                    self._conn.unregister("df_view")
+                page_num += 1
+                # 导入数据到 duckdb
+                self._conn.append(table_name, df)
+                total_count += len(df)
+                db_conn.logger.info(
+                    f"Page {page_num}: imported {total_count} rows to DuckDB table {table_name}"
+                )
+
+        db_conn.logger.info(
+            f"Import completed: {total_count} rows imported to DuckDB table {table_name}"
+        )
+
+    @staticmethod
+    def to_list(result_list: list, description: list, default_value: Any = None) -> list[dict]:
+        """
+        将查询结果转换为字典列表
+
+        Args:
+            result_list: 查询结果
+            description: 列描述
+            default_value: 默认值
+
+        Returns:
+            字典列表
+        """
+        columns = [col[0] for col in description]
+        res = []
+        for data_item in result_list:
+            res_item = {}
+            for i, col_name in enumerate(columns):
+                res_item[col_name] = data_item[i] if data_item[i] is not None else default_value
+            res.append(res_item)
+        return res
+
+    def __del__(self):
+        self.close()