playwrite/db/mysql_connection.py

"""
MySQL 数据库连接组件

提供 MySQL 数据库连接和查询接口
"""

import mysql.connector
from mysql.connector import Error
from typing import List, Dict, Any, Optional
from db.base_connection import BaseDatabaseConnection


class MySQLConnection(BaseDatabaseConnection):
    """MySQL 数据库连接类"""

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        初始化数据库连接

        Args:
            config: 数据库配置字典
                - host: 服务器地址
                - port: 端口号（默认 3306）
                - database: 数据库名称
                - username: 用户名
                - password: 密码
                - charset: 字符集（默认 utf8mb4）
        """
        super().__init__(config)

    def connect(self):
        """
        建立数据库连接

        Returns:
            mysql.connector.connection.MySQLConnection: 数据库连接对象
        """
        if self.connection is not None:
            return self.connection

        try:
            self.connection = mysql.connector.connect(
                host=self.config.get("host", "localhost"),
                port=self.config.get("port", 3306),
                database=self.config["database"],
                user=self.config["username"],
                password=self.config["password"],
                charset=self.config.get("charset", "utf8mb4"),
                autocommit=False,
            )
            print(
                f"成功连接到 MySQL 数据库: {self.config.get('host', 'localhost')}"
                f":{self.config.get('port', 3306)}/{self.config['database']}"
            )
            return self.connection
        except Error as e:
            print(f"MySQL 数据库连接失败: {e}")
            raise

    def disconnect(self):
        """关闭数据库连接"""
        if self.connection and self.connection.is_connected():
            self.connection.close()
            self.connection = None
            print("MySQL 数据库连接已关闭")

    def execute_query(
        self, sql: str, params: Optional[tuple] = None
    ) -> List[Dict[str, Any]]:
        """
        执行查询语句并返回结果

        Args:
            sql: SQL 查询语句
            params: 查询参数（可选）

        Returns:
            List[Dict[str, Any]]: 查询结果列表，每个元素为一行数据的字典
        """
        if not self.connection or not self.connection.is_connected():
            self.connect()

        cursor = None
        try:
            cursor = self.connection.cursor(dictionary=True)
            # Convert SQL Server placeholders (?) to MySQL placeholders (%s)
            # Convert SQL Server table names to MySQL format
            converted_sql = self._convert_placeholders(sql)
            converted_sql = self._convert_table_names(converted_sql)
            if params:
                cursor.execute(converted_sql, params)
            else:
                cursor.execute(converted_sql)

            # 直接获取字典列表
            results = cursor.fetchall()
            return results

        except Error as e:
            print(f"查询执行失败: {e}")
            raise
        finally:
            if cursor:
                cursor.close()

    def execute_update(self, sql: str, params: Optional[tuple] = None) -> int:
        """
        执行更新/插入/删除语句

        Args:
            sql: SQL 语句
            params: 参数（可选）

        Returns:
            int: 受影响的行数
        """
        if not self.connection or not self.connection.is_connected():
            self.connect()

        cursor = None
        try:
            cursor = self.connection.cursor()
            # Convert SQL Server placeholders (?) to MySQL placeholders (%s)
            # Convert SQL Server table names to MySQL format
            converted_sql = self._convert_placeholders(sql)
            converted_sql = self._convert_table_names(converted_sql)
            if params:
                cursor.execute(converted_sql, params)
            else:
                cursor.execute(converted_sql)

            self.connection.commit()
            return cursor.rowcount

        except Error as e:
            self.connection.rollback()
            print(f"执行失败，已回滚: {e}")
            raise
        finally:
            if cursor:
                cursor.close()

    def _convert_placeholders(self, sql: str) -> str:
        """
        Convert SQL Server placeholders (?) to MySQL placeholders (%s)

        This is necessary because the codebase was originally designed for SQL Server,
        which uses '?' as parameter placeholders. MySQL uses '%s' instead.

        Args:
            sql: SQL query with potential SQL Server placeholders

        Returns:
            str: SQL query with MySQL-compatible placeholders
        """
        return sql.replace("?", "%s")

    def _convert_table_names(self, sql: str) -> str:
        """
        Convert SQL Server table names to MySQL format

        Converts [dbo].[TableName] to dbo_TableName and removes square brackets
        from column names.

        Args:
            sql: SQL query with SQL Server table/column names

        Returns:
            str: SQL query with MySQL-compatible table/column names
        """
        import re

        # Convert [dbo].[TableName] to dbo_TableName
        sql = re.sub(r"\[dbo\]\.\[([^\]]+)\]", r"dbo_\1", sql)

        # Remove square brackets from column names (e.g., [Column] -> Column)
        sql = re.sub(r"\[([^\]]+)\]", r"\1", sql)

        return sql

    def get_placeholder(self) -> str:
        """
        获取参数占位符

        Returns:
            MySQL 使用 "%s" 作为参数占位符
        """
        return "%s"