目标:

  1. 连接池管理:自动管理数据库连接,提高性能
  2. 事务支持:通过上下文管理器支持事务操作
  3. 类型提示:完整的类型注解,便于IDE提示
  4. 错误处理:完善的异常处理和日志记录
  5. 多种操作方式:
  6. 原生SQL执行
  7. 便捷的CRUD方法
  8. 批量操作支持
  9. pandas DataFrame集成

一 . 安装以下依赖:

pip install pymysql pandas

二 . 工具类代码:

import pymysql
from typing import List, Tuple,  Dict, Optional
import logging
from contextlib import contextmanager
import pandas as pd

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class MySQLUtil:
    """
    MySQL 数据库操作工具类
    支持连接池、事务管理、批量操作等功能
    """

    def __init__(self, host: str, user: str, password: str, database: str,
                 port: int = 3306, charset: str = 'utf8mb4',
                 pool_size: int = 5, **kwargs):
        """
        初始化数据库连接参数

        Args:
            host: 数据库主机
            user: 用户名
            password: 密码
            database: 数据库名
            port: 端口,默认3306
            charset: 字符集,默认utf8mb4
            pool_size: 连接池大小
            **kwargs: 其他连接参数
        """
        self.db_config = {
            'host': host,
            'user': user,
            'password': password,
            'database': database,
            'port': port,
            'charset': charset,
            'autocommit': False,  # 默认不自动提交,便于事务控制
            **kwargs
        }
        self.pool_size = pool_size
        self._connection_pool = []
        self._initialize_pool()

    def _initialize_pool(self):
        """初始化连接池"""
        try:
            for _ in range(self.pool_size):
                conn = pymysql.connect(**self.db_config)
                self._connection_pool.append(conn)
            logger.info(f"数据库连接池初始化成功,大小: {self.pool_size}")
        except Exception as e:
            logger.error(f"数据库连接池初始化失败: {e}")
            raise

    @contextmanager
    def get_connection(self):
        """
        获取数据库连接(上下文管理器)

        Yields:
            connection: 数据库连接对象
        """
        conn = None
        try:
            if self._connection_pool:
                conn = self._connection_pool.pop()
            else:
                conn = pymysql.connect(**self.db_config)

            yield conn
        except Exception as e:
            if conn:
                conn.rollback()
            logger.error(f"数据库操作失败: {e}")
            raise
        finally:
            if conn:
                # 如果连接池未满,则放回连接池,否则关闭连接
                if len(self._connection_pool) < self.pool_size:
                    self._connection_pool.append(conn)
                else:
                    conn.close()

    def execute_query(self, sql: str, params: Tuple = None) -> List[Tuple]:
        """
        执行查询语句,返回所有结果

        Args:
            sql: SQL查询语句
            params: 参数元组

        Returns:
            查询结果列表
        """
        with self.get_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(sql, params or ())
                result = cursor.fetchall()
                logger.info(f"查询执行成功,影响行数: {cursor.rowcount}")
                return result

    def execute_query_one(self, sql: str, params: Tuple = None) -> Optional[Tuple]:
        """
        执行查询语句,返回单条结果

        Args:
            sql: SQL查询语句
            params: 参数元组

        Returns:
            单条查询结果
        """
        with self.get_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(sql, params or ())
                result = cursor.fetchone()
                logger.info(f"查询执行成功,影响行数: {cursor.rowcount}")
                return result

    def execute_update(self, sql: str, params: Tuple = None) -> int:
        """
        执行更新操作(INSERT, UPDATE, DELETE)

        Args:
            sql: SQL语句
            params: 参数元组

        Returns:
            影响的行数
        """
        with self.get_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(sql, params or ())
                conn.commit()
                affected_rows = cursor.rowcount
                logger.info(f"更新操作执行成功,影响行数: {affected_rows}")
                return affected_rows

    def execute_many(self, sql: str, params_list: List[Tuple]) -> int:
        """
        批量执行操作

        Args:
            sql: SQL语句
            params_list: 参数列表

        Returns:
            影响的总行数
        """
        with self.get_connection() as conn:
            with conn.cursor() as cursor:
                cursor.executemany(sql, params_list)
                conn.commit()
                affected_rows = cursor.rowcount
                logger.info(f"批量操作执行成功,影响行数: {affected_rows}")
                return affected_rows

    def query_to_dataframe(self, sql: str, params: Tuple = None) -> pd.DataFrame:
        """
        执行查询并将结果转换为 pandas DataFrame

        Args:
            sql: SQL查询语句
            params: 参数元组

        Returns:
            pandas DataFrame
        """
        with self.get_connection() as conn:
            df = pd.read_sql(sql, conn, params=params)
            logger.info(f"DataFrame查询成功,返回 {len(df)} 行数据")
            return df

    @contextmanager
    def transaction(self):
        """
        事务上下文管理器

        Usage:
            with db.transaction() as conn:
                db.execute_update("UPDATE ...", conn=conn)
                db.execute_update("INSERT ...", conn=conn)
        """
        conn = None
        try:
            # 从事务中获取新连接,不使用连接池
            conn = pymysql.connect(**self.db_config)
            yield conn
            conn.commit()
            logger.info("事务提交成功")
        except Exception as e:
            if conn:
                conn.rollback()
                logger.error("事务回滚")
            raise
        finally:
            if conn:
                conn.close()

    def close_pool(self):
        """关闭连接池中的所有连接"""
        for conn in self._connection_pool:
            try:
                conn.close()
            except:
                pass
        self._connection_pool.clear()
        logger.info("数据库连接池已关闭")

    def __del__(self):
        """析构函数,确保连接池被正确关闭"""
        self.close_pool()


# 使用示例和测试代码
if __name__ == "__main__":
    # 数据库配置
    DB_CONFIG = {
        'host': 'localhost',
        'user': 'your_username',
        'password': 'your_password',
        'database': 'test_db',
        'port': 3306
    }

    # 创建数据库工具实例
    db_util = MySQLUtil(**DB_CONFIG)

    try:
        # # 1. 创建测试表
        # create_table_sql = """
        #         CREATE TABLE IF NOT EXISTS users (
        #             id INT AUTO_INCREMENT PRIMARY KEY,
        #             name VARCHAR(100) NOT NULL,
        #             email VARCHAR(100) UNIQUE NOT NULL,
        #             age INT,
        #             created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        #         )
        #         """
        # db_util.execute_update(create_table_sql)
        # print("测试表创建成功")

        # 2. 插入单条数据
        user_data = {
            'name': '张三',
            'email': 'zhangsan@example.com',
            'age': 25
        }
        db_util.insert('users', user_data)
        print("单条数据插入成功")

        # 3. 批量插入数据
        users_data = [
            {'name': '李四', 'email': 'lisi@example.com', 'age': 30},
            {'name': '王五', 'email': 'wangwu@example.com', 'age': 28},
            {'name': '赵六', 'email': 'zhaoliu@example.com', 'age': 35}
        ]
        db_util.batch_insert('users', users_data)
        print("批量数据插入成功")

        # 4. 查询数据
        users = db_util.select('users', where='age > %s', where_params=(25,))
        print("查询结果:")
        for user in users:
            print(user)

        # 5. 更新数据
        db_util.update('users', {'age': 26}, 'name = %s', ('张三',))
        print("数据更新成功")

        # 6. 查询到DataFrame
        df = db_util.query_to_dataframe("SELECT * FROM users")
        print("DataFrame数据:")
        print(df)

        # 7. 事务操作示例
        with db_util.transaction() as conn:
            # 在事务中插入数据
            new_user = {'name': '事务用户', 'email': 'transaction@example.com', 'age': 40}
            db_util.insert('users', new_user, conn=conn)

            # 更新年龄
            db_util.update('users', {'age': 41}, 'email = %s', ('transaction@example.com',), conn=conn)
            print("事务操作完成")


    except Exception as e:
        print(f"操作失败: {e}")
    finally:
        # 清理测试数据
        db_util.execute_update("DROP TABLE IF EXISTS users")
        db_util.close_pool()

三 . 方法

  • execute_query(): 执行查询,返回所有结果

  • execute_query_one(): 执行查询,返回单条结果

  • execute_update(): 执行增删改操作

  • execute_many(): 批量执行操作

  • query_to_dataframe(): 查询结果转为DataFrame

  • insert()/batch_insert(): 插入数据

使用时只需修改数据库配置即可快速开始使用。

四 . 画图说明

成功
失败
开始使用 MySQLUtil
初始化数据库连接
配置数据库参数
host, user, password, database, port
选择操作类型
单条记录操作
批量操作
查询操作
事务操作
插入单条数据
insert
更新单条数据
update
删除单条数据
delete
批量插入
batch_insert
批量执行
execute_many
查询多条
execute_query
查询单条
execute_query_one
转DataFrame
query_to_dataframe
开始事务
transaction
在事务中执行多个操作
执行结果
提交事务
回滚事务
获取操作结果
是否需要更多操作
关闭连接池
close_pool
结束
Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐