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

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.2,DBUtiles版本为3.0.3,其它版本可能不适用

文章编写于2023年12月