Yb-QaAutomationHub/tests/core/db_helper.py

"""MySQL 数据库操作工具 — 测试中的数据库校验/数据修改。

连接配置在 tests/config/test_config.yaml 的 database 段维护。

Usage:
    from tests.core.db_helper import update_risk_control_status

    affected = update_risk_control_status("RHJC040026071600021457570057", status=1)
"""

import pymysql

from contextlib import contextmanager

from tests.core.config_loader import load_config


def _resolve_config(config: dict | None) -> dict:
    """显式传入优先，否则读 yaml 的 database 段。"""
    if config is not None:
        return config
    cfg = dict(load_config()["database"])
    cfg.setdefault("charset", "utf8mb4")
    return cfg


def get_connection(config: dict | None = None):
    """获取数据库连接。"""
    return pymysql.connect(**_resolve_config(config))


def execute_sql(sql: str, params: tuple | None = None, config: dict | None = None) -> int:
    """执行 SQL（INSERT/UPDATE/DELETE），返回影响行数。"""
    conn = get_connection(config)
    try:
        with conn.cursor() as cursor:
            affected = cursor.execute(sql, params)
            conn.commit()
            return affected
    finally:
        conn.close()


def query_sql(sql: str, params: tuple | None = None, config: dict | None = None) -> list[dict]:
    """执行查询 SQL，返回结果列表。"""
    conn = get_connection(config)
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            cursor.execute(sql, params)
            return cursor.fetchall()
    finally:
        conn.close()


def update_risk_control_status(order_code: str, status: int = 1, config: dict | None = None) -> int:
    """更新运单风控状态。

    执行 SQL:
        UPDATE tms_order_detail_ext SET risk_control_status = <status>
        WHERE order_detail_code IN (
            SELECT `code` FROM `tms_order_detail` WHERE order_code = '<order_code>'
        )

    Args:
        order_code: 运输单号（对应 tms_order_detail.order_code）
        status: 风控状态值，默认 1
        config: 数据库连接配置（不传则读 yaml database 段）
    Returns:
        影响行数
    """
    sql = (
        "UPDATE tms_order_detail_ext SET risk_control_status = %s "
        "WHERE order_detail_code IN ("
        "  SELECT `code` FROM `tms_order_detail` WHERE order_code = %s"
        ")"
    )
    return execute_sql(sql, (status, order_code), config)


class Db:
    """长连接数据库客户端 — 用例/pipeline 级复用。

    与上面的模块级函数（每语句新建连接）不同，本类在一个实例内复用单条连接：
        - execute/query 前自动 ping 重连，规避长空闲连接被服务端断开
        - 事务外 execute 保持逐语句提交（与旧行为一致）
        - transaction() 内的 execute 不再单独提交，由上下文统一 commit/rollback
        - mogrify() 渲染参数化 SQL 为完整语句（dry-run 预览用，不执行）

    Usage:
        db = Db()                       # 不传 config 则读 yaml database 段
        affected = db.execute("UPDATE t SET a = %s WHERE code = %s", (1, code))
        rows = db.query("SELECT * FROM t WHERE code = %s", (code,))
        with db.transaction():
            db.execute(sql1)
            db.execute(sql2)
        db.close()
    """

    def __init__(self, config: dict | None = None):
        self._kwargs = _resolve_config(config)
        self._conn: pymysql.connections.Connection | None = None
        self._in_transaction = False

    def _connection(self):
        if self._conn is None:
            self._conn = pymysql.connect(**self._kwargs)
        else:
            self._conn.ping(reconnect=True)
        return self._conn

    def execute(self, sql: str, params: tuple | None = None) -> int:
        """执行写语句，返回影响行数。事务内调用不单独提交。"""
        conn = self._connection()
        with conn.cursor() as cursor:
            affected = cursor.execute(sql, params)
        if not self._in_transaction:
            conn.commit()
        return affected

    def query(self, sql: str, params: tuple | None = None) -> list[dict]:
        """执行查询，返回 DictCursor 结果列表。"""
        conn = self._connection()
        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            cursor.execute(sql, params)
            return cursor.fetchall()

    @contextmanager
    def transaction(self):
        """事务上下文：块内所有 execute 在退出时统一提交，异常统一回滚。"""
        conn = self._connection()
        self._in_transaction = True
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            self._in_transaction = False

    def mogrify(self, sql: str, params: tuple | None = None) -> str:
        """渲染参数化 SQL 为最终语句（含引号转义），仅预览不执行。"""
        conn = self._connection()
        with conn.cursor() as cursor:
            return cursor.mogrify(sql, params)

    def close(self) -> None:
        """关闭连接（幂等，可重复调用）。"""
        if self._conn is not None:
            self._conn.close()
            self._conn = None