dataops-sdk 0.5.0__tar.gz → 0.5.1__tar.gz

{dataops_sdk-0.5.0 → dataops_sdk-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataops-sdk
-Version: 0.5.0
+Version: 0.5.1
 Summary: Python SDK for the DataOps platform.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
@@ -16,8 +16,11 @@ Provides-Extra: runner
 Requires-Dist: paramiko>4; extra == "runner"
 Requires-Dist: pandas; extra == "runner"
 Requires-Dist: sqlalchemy; extra == "runner"
-Requires-Dist: pymysql; extra == "runner"
-Requires-Dist: doris-python>=1.0.3; extra == "runner"
+Requires-Dist: mariadb[binary]; extra == "runner"
+Requires-Dist: psycopg[binary]>=3; extra == "runner"
+Requires-Dist: dmpython~=2.5.32; platform_system != "Darwin" and extra == "runner"
+Requires-Dist: dmsqlalchemy<3,>2; platform_system != "Darwin" and extra == "runner"
+Requires-Dist: doris-python>=1.0.4; extra == "runner"
 Requires-Dist: oracledb~=4.0.1; extra == "runner"
 Requires-Dist: neo4j; extra == "runner"
 Requires-Dist: kafka-python>3; extra == "runner"

{dataops_sdk-0.5.0 → dataops_sdk-0.5.1}/pyproject.toml

@@ -27,8 +27,11 @@ runner = [
     "pandas",
     # 关系数据库
     "sqlalchemy",
-    "pymysql",
-    "doris-python>=1.0.3",
+    "mariadb[binary]",
+    "psycopg[binary]>=3",
+    "dmpython~=2.5.32; platform_system != 'Darwin'",
+    "dmsqlalchemy>2,<3; platform_system != 'Darwin'",
+    "doris-python>=1.0.4",
     "oracledb~=4.0.1",
     # 图数据库
     "neo4j",

{dataops_sdk-0.5.0 → dataops_sdk-0.5.1}/src/dataops_sdk/__init__.py

@@ -6,7 +6,7 @@ Nacos 集成能力，未安装时仍可独立导入本包核心功能。
 """
 
 from . import runner, schemas, utils
-from .runner import service_get_datasource_instent
+from .runner import service_get_datasource_instent, get_service_instance
 from .services import (  # noqa: F401
     release_list_create,
     release_list_get_by_code,
@@ -111,4 +111,4 @@ from .services import (  # noqa: F401
 )
 from .session import dataops_base_url_context, request_context
 
-__version__ = "0.5.0"
+__version__ = "0.5.1"

{dataops_sdk-0.5.0 → dataops_sdk-0.5.1}/src/dataops_sdk/runner/__init__.py

@@ -15,6 +15,7 @@ from .datasource_factory import service_get_datasource_instent
 from .schemas.nacos import NacosSettings
 from .utils import (
     DatabaseConn,
+    DatabaseConnAsync,
     FTPFileTransfer,
     KafkaQueue,
     Neo4jConn,
@@ -30,6 +31,7 @@ from .utils import (
 
 __all__ = [
     "DatabaseConn",
+    "DatabaseConnAsync",
     "FTPFileTransfer",
     "KafkaQueue",
     "SFTPFileTransfer",

{dataops_sdk-0.5.0 → dataops_sdk-0.5.1}/src/dataops_sdk/runner/datasource_factory.py

@@ -6,24 +6,33 @@
 from dataops_sdk.services.datasource import service_get_datasource_by_name
 
 
-def service_get_datasource_instent(datasource: str, log_sql: bool | None = None, **kwargs):
-    """
-    根据数据源名称获取对应的连接实例。
+def service_get_datasource_instent(
+    datasource: str,
+    log_sql: bool | None = None,
+    use_async: bool = False,
+    **kwargs,
+):
+    """根据数据源名称获取对应的连接实例.
 
     通过数据源名称从平台查询元数据，根据数据源的 category（数据库/消息队列）
     和 type（具体类型）动态创建并返回相应的连接实例。
 
+    ``use_async=True`` 时，关系型数据库返回 ``DatabaseConnAsync``。
+    Neo4j、Kafka 以及达梦（dm）尚无异步实现。
+
     Args:
         datasource: 数据源名称，用于从平台查询数据源元数据。
         log_sql: 是否记录 SQL 执行日志，传递给数据库连接实例。
+        use_async: 为 True 时返回异步连接。默认为 False，保持同步连接。
+        **kwargs: 预留的额外参数，当前不会传给连接实例。
 
     Returns:
-        DatabaseConn: 通用数据库连接实例（MySQL、PostgreSQL 等）。
+        DatabaseConn | DatabaseConnAsync: 通用数据库连接实例（MySQL、PostgreSQL 等）。
         Neo4jConn: Neo4j 图数据库连接实例。
         KafkaQueue: Kafka 消息队列连接实例。
 
     Raises:
-        NotImplementedError: 当数据源类型未被支持时抛出。
+        NotImplementedError: 当数据源类型未被支持，或该类型没有异步实现时抛出。
 
     Note:
         当前支持的 category:
@@ -39,6 +48,8 @@ def service_get_datasource_instent(datasource: str, log_sql: bool | None = None,
         if meta.data.type == "kafka":
             # Kafka 连接需要 bootstrap_servers
             assert meta.data.bootstrap_servers, "bootstrap_servers is required, but not provided"
+            if use_async:
+                raise NotImplementedError("kafka does not support async connection")
             from dataops_sdk.runner import KafkaQueue
 
             return KafkaQueue(meta.data.bootstrap_servers)
@@ -54,6 +65,8 @@ def service_get_datasource_instent(datasource: str, log_sql: bool | None = None,
         assert meta.data.database, "database is required, but not provided"
 
         if meta.data.type == "neo4j":
+            if use_async:
+                raise NotImplementedError("Neo4j does not support async connection")
             # Neo4j 图数据库使用专用连接类
             from dataops_sdk.runner import Neo4jConn
 
@@ -68,6 +81,19 @@ def service_get_datasource_instent(datasource: str, log_sql: bool | None = None,
             )
         else:
             # 通用数据库连接，支持 MySQL、PostgreSQL 等
+            if use_async:
+                from dataops_sdk.runner import DatabaseConnAsync
+
+                return DatabaseConnAsync(
+                    host=meta.data.ip,
+                    port=meta.data.port,
+                    user=meta.data.username,
+                    password=meta.data.password,
+                    db=meta.data.database,
+                    db_type=meta.data.type,
+                    log_sql=log_sql,
+                    name=meta.title,
+                )
             from dataops_sdk.runner import DatabaseConn
 
             return DatabaseConn(

{dataops_sdk-0.5.0 → dataops_sdk-0.5.1}/src/dataops_sdk/runner/utils/__init__.py

@@ -10,6 +10,7 @@ Functions:
 
 Classes:
     DatabaseConn: 数据库连接管理类.
+    DatabaseConnAsync: 异步数据库连接管理类.
     Neo4jConn: Neo4j 图数据库连接类.
     FTPFileTransfer: FTP 文件传输类.
     SFTPFileTransfer: SFTP 文件传输类.
@@ -27,6 +28,7 @@ from .common import (
     parse_trigger_time,
     print_json,
 )
+from .dbConn_async import DatabaseConnAsync
 from .dbConn_sync import DatabaseConn
 from .ftp_file_transfer import FTPFileTransfer
 from .nacos import get_service_instance
@@ -41,6 +43,7 @@ __all__ = [
     "generate_uuid",
     "print_json",
     "DatabaseConn",
+    "DatabaseConnAsync",
     "Neo4jConn",
     "get_service_instance",
     "KafkaQueue",

dataops_sdk-0.5.1/src/dataops_sdk/runner/utils/dbConn_async.py

@@ -0,0 +1,484 @@
+"""基于 SQLAlchemy 的异步数据库连接实现.
+
+继承自 DatabaseConnBaseAsync，提供与 DatabaseConn 对齐的异步接口。
+使用 SQLAlchemy AsyncEngine 和 AsyncSession 作为底层实现。
+
+Classes:
+    DatabaseConnAsync: 基于 SQLAlchemy 的异步数据库连接类.
+    DatabaseConnDataOpsAsync: 通过数据源名称动态创建的 DatabaseConnAsync 子类.
+
+Example:
+    >>> async with DatabaseConnAsync(host="localhost", port=3306,
+    ...                              user="root", password="pass", db="test") as conn:
+    ...     result = await conn.select("SELECT * FROM users WHERE id = %s", (1,))
+    ...     await conn.execute("INSERT INTO users (name) VALUES (%s)", ("Alice",))
+"""
+
+import re
+from collections.abc import Mapping
+from typing import Any, Final, Literal, overload
+
+from sqlalchemy import URL, text
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+
+from .dbConn_base import DatabaseConnBaseAsync
+
+# 数据库类型到异步驱动程序的映射
+ASYNC_DBTYPE_DRIVER: Final[Mapping[str, str]] = {
+    "mysql": "mariadbconnector",
+    "mariadb": "mariadbconnector",
+    "postgresql": "psycopg",
+    "oracle": "oracledb",
+    "doris": "mariadbconnector",
+}
+
+
+def build_async_url(
+    host: str,
+    port: int | str,
+    user: str,
+    password: str,
+    db: str,
+    db_type: str,
+) -> URL:
+    """构建异步数据库连接 URL.
+
+    Args:
+        host: 数据库服务器主机地址.
+        port: 数据库服务器端口号.
+        user: 数据库用户名.
+        password: 数据库密码.
+        db: 数据库名称.
+        db_type: 数据库类型.
+
+    Returns:
+        URL: SQLAlchemy 异步连接 URL.
+
+    Raises:
+        NotImplementedError: 当指定的 db_type 没有异步驱动时抛出.
+    """
+    if db_type not in ASYNC_DBTYPE_DRIVER:
+        raise NotImplementedError(f"{db_type} 异步驱动尚未支持，请联系管理员")
+
+    query: dict[str, str] = {}
+    if db_type in {"mysql", "mariadb", "doris"}:
+        query["local_infile"] = "True"
+        query["ssl_disabled"] = "True"
+
+    driver = ASYNC_DBTYPE_DRIVER[db_type]
+    return URL.create(
+        drivername=f"{db_type}+{driver}",
+        username=user,
+        password=password,
+        host=host,
+        port=int(port),
+        database=db,
+        query=query,
+    )
+
+
+class DatabaseConnAsync(DatabaseConnBaseAsync):
+    """基于 SQLAlchemy 的异步数据库连接类.
+
+    使用 AsyncEngine 和 AsyncSession 实现数据库操作，支持单条执行、批量执行、
+    查询、手动事务以及异步上下文管理器。
+
+    Attributes:
+        engine: SQLAlchemy AsyncEngine 对象.
+        session: 当前事务的 AsyncSession，未开启事务时为 None.
+        session_factory: SQLAlchemy 异步 Session 工厂.
+        logger: 日志记录器.
+
+    Example:
+        >>> async with DatabaseConnAsync(
+        ...     host="localhost", port=3306, user="root", password="pass",
+        ...     db="test", db_type="mysql",
+        ... ) as conn:
+        ...     rows = await conn.select("SELECT * FROM users LIMIT %s", (10,))
+        ...     await conn.execute(
+        ...         "INSERT INTO users (name) VALUES (%s)", ("Alice",)
+        ...     )
+    """
+
+    def post_init(self):
+        """初始化 SQLAlchemy 异步 Engine 和 Session 工厂.
+
+        Returns:
+            None
+        """
+        self._url = build_async_url(
+            host=self._host,
+            port=self._port,
+            user=self._user,
+            password=self._password,
+            db=self._db,
+            db_type=self._db_type,
+        )
+
+        if self._log_sql:
+            self._logger.info(
+                f"Creating SQLAlchemy async Engine for {self._host}:{self._port}/{self._db}"
+            )
+        self._engine: AsyncEngine | None = create_async_engine_safe(self._db_type, self._url)
+        self._session_factory = async_sessionmaker(
+            bind=self._engine,
+            autoflush=False,
+            expire_on_commit=False,
+        )
+        self._session: AsyncSession | None = 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 语句，仅用于日志展示.
+
+        Args:
+            sql: SQL 语句模板.
+            args: SQL 语句参数.
+
+        Returns:
+            str: 渲染后的 SQL 语句.
+        """
+        if not args:
+            return sql
+        if isinstance(args, (list, tuple)):
+            return sql.replace("%s", "{}").format(*args)
+        sql = re.sub(r"\%\((\w+)\)s", r"\{\1\}", sql)
+        return sql.format(**args)
+
+    async def executemany(
+        self,
+        sql: str,
+        datas: list[tuple | list],
+        *,
+        auto_commit: bool = True,
+        log_sql: bool | None = None,
+        **kwargs,
+    ):
+        """执行多条 SQL 语句（批量执行）.
+
+        已通过 begin() 开启手动事务时，使用该事务执行，auto_commit 被忽略。
+
+        Args:
+            sql: 待执行的 SQL 语句模板，支持 %s 或 :name 占位符.
+            datas: SQL 语句参数的列表，每一项对应一条 SQL 执行.
+            auto_commit: 是否自动提交事务，默认为 True. 存在活跃事务时此参数被忽略.
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql.
+            **kwargs: 兼容性的额外参数，传入时会输出警告日志.
+        """
+        if kwargs:
+            self.logger.warning(f"尚未支持的参数: {kwargs}")
+
+        should_log = self._get_log_sql(log_sql)
+        sql, args = self.process_input(sql, datas)
+
+        if should_log and args:
+            self.logger.info(self.render_sql(sql, args[0]))
+
+        if self._session is not None:
+            await self._session.execute(text(sql), args)
+            return
+
+        async with self.session_factory() as session:
+            await session.execute(text(sql), args)
+            if auto_commit:
+                await session.commit()
+
+    async def execute(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        *,
+        auto_commit: bool = True,
+        log_sql: bool | None = None,
+        **kwargs,
+    ):
+        """执行单条 SQL 语句.
+
+        已通过 begin() 开启手动事务时，使用该事务执行，auto_commit 被忽略。
+
+        Args:
+            sql: 待执行的 SQL 语句.
+            args: SQL 语句参数，支持元组、列表或字典格式.
+            auto_commit: 是否自动提交事务，默认为 True. 存在活跃事务时此参数被忽略.
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql.
+            **kwargs: 兼容性的额外参数，传入时会输出警告日志.
+        """
+        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 is not None:
+            await self._session.execute(text(sql), args)
+            return
+
+        async with self.session_factory() as session:
+            await session.execute(text(sql), args)
+            if auto_commit:
+                await session.commit()
+
+    @overload
+    async 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, ...]]: ...
+
+    @overload
+    async def select(
+        self,
+        sql: str,
+        args: tuple | list | dict | None = None,
+        size: int | None = None,
+        mapping: Literal[True] = ...,
+        *,
+        log_sql: bool | None = None,
+    ) -> list[dict[str, Any]]: ...
+
+    async 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, ...]]:
+        """执行查询语句并返回结果.
+
+        Args:
+            sql: 待执行的查询 SQL 语句.
+            args: SQL 语句参数，支持元组、列表或字典格式.
+            size: 限制返回的行数，None 或小于等于 0 时不限制.
+            mapping: 是否以字典形式返回每行（列名 -> 值），默认为 False.
+            log_sql: 是否打印 SQL 日志，为 None 时使用实例级 self._log_sql.
+
+        Returns:
+            list[dict[str, Any] | tuple[Any, ...]]: 查询结果列表；
+            当 mapping=True 时元素为字典，否则元素为元组.
+        """
+        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))
+
+        async with self.session_factory() as session:
+            connection = await session.connection()
+            if isinstance(sql, str):
+                result = await connection.exec_driver_sql(sql, args)
+            else:
+                result = await connection.execute(sql)
+
+            if mapping:
+                result = result.mappings()
+            rows = result.fetchall() if size is None or size <= 0 else result.fetchmany(size)
+
+        if mapping:
+            return [dict(row) for row in rows]
+        return [row._t._data for row in rows]
+
+    @staticmethod
+    def process_input(sql: str, args: dict | tuple | list | None):
+        """预处理 SQL 参数.
+
+        将列表形式的参数规范化为 SQLAlchemy 接受的元组形式。
+
+        Args:
+            sql: SQL 语句模板.
+            args: SQL 参数，支持 None、元组、列表或字典.
+
+        Returns:
+            tuple[str, Any]: 处理后的 (sql, args) 元组.
+        """
+        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))
+            return sql, tuple(args)
+        return sql, args
+
+    async def close(self):
+        """关闭数据库连接和活跃事务.
+
+        如果存在未提交的事务，先执行回滚，然后关闭 Session 和 Engine。
+        支持多次调用。
+        """
+        if getattr(self, "_session", None) is not None:
+            try:
+                await self._session.rollback()
+            except Exception:
+                pass
+            finally:
+                await self._session.close()
+                self._session = None
+
+        if getattr(self, "_engine", None) is not None:
+            await self._engine.dispose()
+            self._engine = None
+            self._logger.info("closed")
+
+    @property
+    def session_factory(self):
+        """获取 SQLAlchemy 异步 Session 工厂.
+
+        Returns:
+            async_sessionmaker: 已绑定 AsyncEngine 的 Session 工厂.
+        """
+        return self._session_factory
+
+    @property
+    def engine(self) -> AsyncEngine | None:
+        """获取 SQLAlchemy AsyncEngine 对象.
+
+        Returns:
+            AsyncEngine | None: 异步 Engine，未初始化时为 None.
+        """
+        return self._engine
+
+    async def commit(self):
+        """提交当前手动事务.
+
+        仅在已通过 begin() 开启事务后调用。提交后事务结束，session 会被关闭。
+        """
+        if self._session is None:
+            self._logger.warning("commit called but no active transaction")
+            return
+        self._logger.debug("committing transaction")
+        await self._session.commit()
+        await self._session.close()
+        self._session = None
+
+    async def rollback(self):
+        """回滚当前手动事务.
+
+        仅在已通过 begin() 开启事务后调用。回滚后事务结束，session 会被关闭。
+        """
+        if self._session is None:
+            self._logger.warning("rollback called but no active transaction")
+            return
+        self._logger.warning("rolling back transaction")
+        await self._session.rollback()
+        await self._session.close()
+        self._session = None
+
+    async def begin(self):
+        """开启一个新的手动事务.
+
+        开启后，后续的 execute() 和 executemany() 将在该事务中执行。
+        调用 commit() 或 rollback() 会结束当前事务。
+
+        Raises:
+            RuntimeError: 如果已有活跃事务未关闭时再次调用。
+        """
+        if self._session is not None:
+            raise RuntimeError(
+                "Transaction already active. Commit or rollback before starting a new one."
+            )
+        self._logger.debug("beginning transaction")
+        self._session = self.session_factory()
+        await self._session.begin()
+
+    @property
+    def session(self) -> AsyncSession | None:
+        """获取当前事务的 Session 对象.
+
+        Returns:
+            AsyncSession | None: 当前事务的 Session；未开启事务时返回 None.
+        """
+        return self._session
+
+    @property
+    def in_transaction(self) -> bool:
+        """检查是否存在活跃的手动事务.
+
+        Returns:
+            bool: 如果存在活跃事务则返回 True，否则返回 False.
+        """
+        return self._session is not None
+
+
+class DatabaseConnDataOpsAsync(DatabaseConnAsync):
+    """基于数据源名称动态创建的 DatabaseConnAsync 子类.
+
+    通过平台元数据服务根据 datasource 名称获取连接信息，
+    自动构造 DatabaseConnAsync 实例。
+
+    Attributes:
+        datasource: 数据源名称.
+        log_sql: 是否记录 SQL 执行日志.
+    """
+
+    def __init__(self, datasource: str, log_sql: bool | None = None):
+        """初始化 DatabaseConnDataOpsAsync.
+
+        Args:
+            datasource: 数据源名称.
+            log_sql: 是否记录 SQL 执行日志，默认为 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__(
+            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_async_engine_safe(db_type: str, url: str | URL, **kwargs: Any) -> AsyncEngine:
+    """安全创建 SQLAlchemy 异步 Engine.
+
+    Args:
+        db_type: 数据库类型，如 mysql、postgresql、oracle、doris.
+        url: 数据库连接 URL.
+        **kwargs: 传递给 create_async_engine 的额外参数.
+
+    Returns:
+        AsyncEngine: 配置好的 SQLAlchemy 异步 Engine 实例.
+
+    Raises:
+        NotImplementedError: 当数据库类型不在异步支持列表中时抛出.
+    """
+    if db_type not in ASYNC_DBTYPE_DRIVER:
+        raise NotImplementedError(f"{db_type} 异步驱动尚未支持，请联系管理员")
+
+    kwargs.setdefault("pool_pre_ping", True)
+    if db_type in ("mysql", "mariadb", "doris"):
+        kwargs.setdefault("pool_recycle", 3600)
+
+    return create_async_engine(url, **kwargs)