120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
|
import aiomysql
|
|||
|
|
|
|||
|
|
class DBPool:
|
|||
|
|
def __init__(self, coin, options):
|
|||
|
|
self.coin = coin
|
|||
|
|
self.pool = None
|
|||
|
|
self.options = options
|
|||
|
|
|
|||
|
|
async def _get_pool(self):
|
|||
|
|
if self.pool is None:
|
|||
|
|
self.pool = await aiomysql.create_pool(**self.options)
|
|||
|
|
return self.pool
|
|||
|
|
|
|||
|
|
# -------------------------
|
|||
|
|
# 工具：判断是否 SELECT
|
|||
|
|
# -------------------------
|
|||
|
|
def _is_select(self, sql: str) -> bool:
|
|||
|
|
return sql.lstrip().lower().startswith("select")
|
|||
|
|
|
|||
|
|
# -------------------------
|
|||
|
|
# 非事务 SQL
|
|||
|
|
# -------------------------
|
|||
|
|
async def exec(self, sql, values=None):
|
|||
|
|
if values is None:
|
|||
|
|
values = []
|
|||
|
|
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
async with conn.cursor(aiomysql.DictCursor) as cur:
|
|||
|
|
await cur.execute(sql, values)
|
|||
|
|
|
|||
|
|
if self._is_select(sql):
|
|||
|
|
# JS: SELECT 返回 rows
|
|||
|
|
return await cur.fetchall()
|
|||
|
|
else:
|
|||
|
|
# JS: 非 SELECT 返回 OkPacket
|
|||
|
|
return {
|
|||
|
|
"affectedRows": cur.rowcount,
|
|||
|
|
"insertId": cur.lastrowid,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# -------------------------
|
|||
|
|
# 单 SQL 事务
|
|||
|
|
# -------------------------
|
|||
|
|
async def exec_transaction(self, sql, values=None):
|
|||
|
|
if values is None:
|
|||
|
|
values = []
|
|||
|
|
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
async with conn.cursor(aiomysql.DictCursor) as cur:
|
|||
|
|
await conn.begin()
|
|||
|
|
try:
|
|||
|
|
await cur.execute(sql, values)
|
|||
|
|
|
|||
|
|
if self._is_select(sql):
|
|||
|
|
result = await cur.fetchall()
|
|||
|
|
else:
|
|||
|
|
result = {
|
|||
|
|
"affectedRows": cur.rowcount,
|
|||
|
|
"insertId": cur.lastrowid,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
await conn.commit()
|
|||
|
|
return result
|
|||
|
|
except Exception:
|
|||
|
|
await conn.rollback()
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# -------------------------
|
|||
|
|
# 多 SQL 合并事务
|
|||
|
|
# -------------------------
|
|||
|
|
async def exec_transaction_together(self, params):
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
async with conn.cursor() as cur:
|
|||
|
|
await conn.begin()
|
|||
|
|
try:
|
|||
|
|
for item in params:
|
|||
|
|
sql = item["sql"]
|
|||
|
|
param = item.get("param", [])
|
|||
|
|
await cur.execute(sql, param)
|
|||
|
|
await conn.commit()
|
|||
|
|
except Exception:
|
|||
|
|
await conn.rollback()
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# -------------------------
|
|||
|
|
# 写锁事务
|
|||
|
|
# -------------------------
|
|||
|
|
async def exec_write_lock(self, sql, values=None, table_name=None):
|
|||
|
|
if values is None:
|
|||
|
|
values = []
|
|||
|
|
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
async with conn.cursor(aiomysql.DictCursor) as cur:
|
|||
|
|
await conn.begin()
|
|||
|
|
try:
|
|||
|
|
await cur.execute(f"LOCK TABLES {table_name} WRITE")
|
|||
|
|
await cur.execute(sql, values)
|
|||
|
|
|
|||
|
|
if self._is_select(sql):
|
|||
|
|
result = await cur.fetchall()
|
|||
|
|
else:
|
|||
|
|
result = {
|
|||
|
|
"affectedRows": cur.rowcount,
|
|||
|
|
"insertId": cur.lastrowid,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
await conn.commit()
|
|||
|
|
return result
|
|||
|
|
except Exception:
|
|||
|
|
await conn.rollback()
|
|||
|
|
raise
|
|||
|
|
finally:
|
|||
|
|
# 对齐 JS：始终解锁
|
|||
|
|
await conn.cursor().execute("UNLOCK TABLES")
|
|||
|
|
|
|||
|
|
__all__ = ["DBPool"]
|