MySQL

类

import pymysql.cursors
from typing import Union, List

class MySQLHelper:
    def __init__(self, host: str, port: int, user: str, password: str, db: str):
        """
        初始化 MySQLHelper 类

        Args:
            host (str): MySQL 主机名
            port (int): MySQL 端口号
            user (str): MySQL 用户名
            password (str): MySQL 密码
            db (str): 数据库名称
        """
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.db = db
        self.connection = None

    def connect(self):
        """
        连接到 MySQL 数据库
        """
        self.connection = pymysql.connect(
            host=self.host,
            port=self.port,
            user=self.user,
            password=self.password,
            db=self.db,
            cursorclass=pymysql.cursors.DictCursor
        )

    def disconnect(self):
        """
        断开与 MySQL 数据库的连接
        """
        if self.connection:
            self.connection.close()

    def execute_query(self, sql: str, params: Union[List, tuple]=None) -> List[dict]:
        """
        执行查询语句

        Args:
            sql (str): 查询语句
            params (Union[List, tuple], optional): 参数列表，默认为 None

        Returns:
            List[dict]: 查询结果
        """
        self.connect()
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(sql, params)
                result = cursor.fetchall()
                return result
        finally:
            self.disconnect()

    def execute_update(self, sql: str, params: Union[List, tuple]=None) -> int:
        """
        执行更新语句

        Args:
            sql (str): 更新语句
            params (Union[List, tuple], optional): 参数列表，默认为 None

        Returns:
            int: 更新的行数
        """
        self.connect()
        try:
            with self.connection.cursor() as cursor:
                result = cursor.execute(sql, params)
                self.connection.commit()
                return result
        finally:
            self.disconnect()

    def insert_single(self, table_name: str, data: Dict[str, Union[str, int, float]]) -> int:
        """
        插入单条数据

        Args:
            table_name (str): 表名
            data (Dict[str, Union[str, int, float]]): 插入的数据字典

        Returns:
            int: 插入的行数
        """
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        values = list(data.values())
        insert_sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
        return self.execute_update(insert_sql, values)

    def insert_multiple(self, table_name: str, data_list: List[Dict[str, Union[str, int, float]]]) -> int:
        """
        批量插入数据

        Args:
            table_name (str): 表名
            data_list (List[Dict[str, Union[str, int, float]]]): 批量插入的数据列表

        Returns:
            int: 插入的总行数
        """
        if not data_list:
            return 0

        columns = ', '.join(data_list[0].keys())
        placeholders = ', '.join(['%s'] * len(data_list[0]))
        values = [tuple(data.values()) for data in data_list]
        insert_sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
        return self.execute_update(insert_sql, values)

使用示例

# 使用示例
# 初始化 MySQLHelper
mysql_helper = MySQLHelper("localhost", 3306, "root", "password", "mydb")

# 执行查询语句
select_sql = "SELECT * FROM users"
result = mysql_helper.execute_query(select_sql)
print("Query result:")
for row in result:
    print(row)

# 执行更新语句
update_sql = "UPDATE users SET age = %s WHERE id = %s"
update_params = (25, 1)
affected_rows = mysql_helper.execute_update(update_sql, update_params)
print(f"Affected rows: {affected_rows}")
# 单条数据插入
single_data = {"name": "John", "age": 25, "email": "john@example.com"}
affected_rows = mysql_helper.insert_single("users", single_data)
print(f"Affected rows (single): {affected_rows}")

# 批量数据插入
data_list = [
    {"name": "Alice", "age": 30, "email": "alice@example.com"},
    {"name": "Bob", "age": 35, "email": "bob@example.com"},
    {"name": "Charlie", "age": 28, "email": "charlie@example.com"}
]
total_rows = mysql_helper.insert_multiple("users", data_list)
print(f"Total affected rows (multiple): {total_rows}")