Pymysql+DBUtils实现Mysql数据库连接池

pymysql并不是线程安全的，所以有多个线程同时使用pymysql操作数据库时就会出现问题。为了线程安全和更高地并发，使用数据库连接池来管理连接。 所以我结合pymysql和DBUtils进行了封装，可以参考使用。

"""
python版本3.8.7
PyMySQL版本1.0.2
"""
import threading
from typing import Literal, Tuple, Optional, Callable, Union

import pymysql
from pymysql.cursors import Cursor
from pymysql.connections import Connection
from dbutils.pooled_db import PooledDB


# 单例模式
class SingletonMeta(type):
   _instances = {}
   _instance_lock = threading.Lock()

   def __call__(cls, *args, **kwargs):
      if cls not in cls._instances:
         with SingletonMeta._instance_lock:
            if cls not in cls._instances:
               cls._instances[cls] = super(SingletonMeta, cls).__call__(*args, **kwargs)
      return cls._instances[cls]


# 异常处理
class MysqlPoolException(Exception):
   ...


# 连接池封装
class MysqlPool(metaclass=SingletonMeta):

   def __init__(self, **kwargs):
      self._pool = PooledDB(**kwargs)

   def _connect(self) -> Tuple[Connection, Cursor]:
      """获取一个连接"""
      conn = self._pool.connection()
      cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
      return conn, cursor

   @staticmethod
   def _close(conn: Connection, cursor: Cursor) -> None:
      """关闭连接"""
      conn.close()
      cursor.close()

   def exec_sql(
         self, callback: Callable,
         *args,
         error: Literal["raise", "output"] = "raise",
         commitable: bool = False,
         rollbackable: bool = False,
         **kwargs
   ):
      """执行sql"""
      conn, cursor = self._connect()
      try:
         ret = callback(*args, **kwargs, cursor=cursor)
      except Exception as e:
         if rollbackable is True:
            conn.rollback()
         if error == "raise":
            raise MysqlPoolException(e) from None
         elif error == "output":
            return str(e)
         else:
            raise ValueError("error参数值只能为'raise'或'output'")
      finally:
         if commitable is True:
            conn.commit()
         self._close(conn, cursor)
      return ret

   @staticmethod
   def _get(sql, params=None, cursor: Cursor = None):
      cursor.execute(sql, params)
      return cursor.fetchone()

   @staticmethod
   def _get_all(sql, params=None, cursor=None):
      cursor.execute(sql, params)
      return cursor.fetchall()

   @staticmethod
   def _post(sql, params=None, cursor: Cursor = None):
      return cursor.execute(sql, params)

   @staticmethod
   def _post_many(sql, params=None, cursor: Cursor = None):
      return cursor.executemany(sql, params)

   def mysql_pool(
         self,
         sql: str,
         params: Optional[list] = None,
         method: Literal["fetchone", "fetchall", "execute", "executemany"] = "fetchall",
         *args,
         **kwargs
   ) -> Union[tuple, list, dict, str, None]:
      mapping = {
         "fetchone": {"callable_": self._get, "commitable": False, "rollbackable": False},
         "fetchall": {"callable_": self._get_all, "commitable": False, "rollbackable": False},
         "execute": {"callable_": self._post, "commitable": True, "rollbackable": False},
         "executemany": {"callable_": self._post_many, "commitable": True, "rollbackable": False},
      }
      return self.exec_sql(sql=sql, params=params, *args, **mapping[method], **kwargs)

   def fetchone(self, sql: str, params: Optional[list] = None, *args, **kwargs) -> Union[dict, str, None]:
      return self.exec_sql(
         callback=self._get,
         sql=sql,
         params=params,
         commitable=False,
         rollbackable=False,
         *args,
         **kwargs
      )

   def fetchall(self, sql: str, params: Optional[list] = None, *args, **kwargs) -> Union[tuple, list, str]:
      """查询一条数据"""
      return self.exec_sql(
         callback=self._get_all,
         sql=sql,
         params=params,
         commitable=False,
         rollbackable=False,
         *args,
         **kwargs
      )

   def execute(self, sql: str, params: Optional[list] = None, *args, **kwargs) -> Union[int, str]:
      """执行sql"""
      return self.exec_sql(
         callback=self._post,
         sql=sql,
         params=params,
         commitable=True,
         rollbackable=False,
         *args,
         **kwargs
      )

   def executemany(self, sql: str, params: Optional[list] = None, *args, **kwargs) -> Union[int, str]:
      """插入多条数据"""
      return self.exec_sql(
         callback=self._post_many,
         sql=sql,
         params=params,
         commitable=True,
         rollbackable=False,
         *args,
         **kwargs
      )


# 连接配置
MYSQL_POOL_CONFIG = {
   "creator": pymysql,
   "mincached": 2,
   "maxshared": 5,
   "maxconnections": 10,
   "blocking": True,
   "maxusage": None,
   "host": "127.0.0.1",
   "port": 3306,
   "user": "root",
   "password": "111111",
   "database": "test_db",
   "charset": "utf8mb4"
}

if __name__ == '__main__':
   import time

   # 自定义回调函数
   def callback(sql, params, cursor: Cursor = None):
      for data in params:
         cursor.execute(sql, data)

   db = MysqlPool(**MYSQL_POOL_CONFIG)

   sql = "insert into user(name,sex,age)values(%s,%s,%s);"
   params = [("张三", "男", 22) for _ in range(10000)]

   start = time.perf_counter()
   # 使用exec_sql方法执行事务操作
   db.exec_sql(callback, sql=sql, params=params, commitable=True, rollbackable=True)
   end = time.perf_counter()
   print("耗时: %.2f" % (end - start))

   start = time.perf_counter()
   # 插入多条数据 速度更快
   db.executemany(sql=sql, params=params)
   end = time.perf_counter()
   print("耗时: %.2f" % (end - start))

注意

我使用的PyMySQL版本为 1.1.0,DBUtiles版本为3.0.3,其它版本可能不适用

文章编写于2023年12月