aiomysql 常用连接方式

适用于 Python 3.11 及以上版本。先安装依赖：python -m pip install aiomysql "PyMySQL[rsa]"。

复制下方代码，保存为 aiomysql_client.py。设置 MYSQL_HOST、MYSQL_USER、MYSQL_PASSWORD、MYSQL_DATABASE，再运行 python aiomysql_client.py；MYSQL_PORT 默认是 3306。

入口使用正文的 intl_sixin_data 示例表，筛选 session_id = 'session_a'，最多返回 5 条记录的 id、date_num、session_id；没有匹配记录时返回 []。建表语句见正文 1.2。连接自己的表时，修改 main() 中的查询。

在其他脚本中可通过 from aiomysql_client import AsyncMySQLClient 导入，使用 await AsyncMySQLClient.create(...) 创建客户端，结束时执行 await client.close()。事务和并发的调用方式见正文。

aiomysql_client.py Python

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
"""aiomysql 异步客户端，适用于 Python 3.11+。

安装依赖：python -m pip install aiomysql "PyMySQL[rsa]"
运行前设置 MYSQL_HOST、MYSQL_USER、MYSQL_PASSWORD、MYSQL_DATABASE；MYSQL_PORT 默认 3306。
直接运行本文件只查询 intl_sixin_data 中的演示会话 session_a；无匹配记录时返回 []。
"""

import asyncio
import os
from contextlib import asynccontextmanager

import aiomysql


class AsyncMySQLClient:
    """使用 create 创建；连接池及其所有操作必须属于同一个事件循环。"""

    def __init__(self, pool):
        self._pool = pool
        self._slots = asyncio.Semaphore(pool.maxsize)

    @classmethod
    async def create(cls, host, user, password, database, port=3306, minsize=1, maxsize=5,
                     connect_timeout=10, pool_recycle=1800):
        if maxsize < 1:
            raise ValueError("maxsize 必须大于 0")
        pool = await aiomysql.create_pool(
            host=host, port=port, user=user, password=password, db=database, charset="utf8mb4",
            cursorclass=aiomysql.DictCursor, autocommit=True, minsize=minsize, maxsize=maxsize,
            connect_timeout=connect_timeout, pool_recycle=pool_recycle,
        )
        return cls(pool)

    @asynccontextmanager
    async def _connection(self):
        # aiomysql 0.3.2 丢弃已关闭连接时不唤醒池等待者，让排队发生在信号量上。
        # 先归还/丢弃连接，再释放名额，下一个任务进入池时即可复用或创建连接。
        async with self._slots:
            async with self._pool.acquire() as conn:
                cursor = await conn.cursor()
                try:
                    yield conn, cursor
                except asyncio.CancelledError:
                    # 未读完的响应不能复用；关闭连接后也不再让游标读取剩余结果集。
                    conn.close()
                    raise
                finally:
                    if not conn.closed:
                        await cursor.close()

    async def query(self, sql: str, params=None) -> list[dict]:
        """执行返回结果集的语句；每次调用独立借用连接。"""
        async with self._connection() as (_, cursor):
            await cursor.execute(sql, params)
            return list(await cursor.fetchall())

    async def execute(self, sql: str, params=None) -> int:
        """在独立事务中执行一条写入语句，返回受影响行数。"""
        async with self.transaction() as cursor:
            return await cursor.execute(sql, params)

    async def executemany(self, sql: str, params) -> int:
        """params 为参数列表；空列表返回 0，非空列表在同一事务中执行。"""
        if not params:
            return 0
        async with self.transaction() as cursor:
            return await cursor.executemany(sql, params)

    @asynccontextmanager
    async def transaction(self):
        """多条语句共用同一游标；正常退出提交，普通异常退出回滚。"""
        async with self._connection() as (conn, cursor):
            try:
                await conn.begin()
                yield cursor
                await conn.commit()
            except Exception:
                try:
                    await conn.rollback()
                except BaseException:
                    # 回滚失败或被取消时，禁止将状态不明的连接放回空闲队列。
                    conn.close()
                    raise
                raise
        # 取消由 _connection 关闭物理连接，服务端回滚尚未提交的事务。

    async def close(self):
        """所有任务结束后调用；停止接收新的借用，并等待连接关闭。"""
        self._pool.close()
        await self._pool.wait_closed()


async def main():
    client = await AsyncMySQLClient.create(
        host=os.environ["MYSQL_HOST"], port=int(os.environ.get("MYSQL_PORT", "3306")),
        user=os.environ["MYSQL_USER"], password=os.environ["MYSQL_PASSWORD"],
        database=os.environ["MYSQL_DATABASE"],
    )
    try:
        sql = ("SELECT id, date_num, session_id FROM intl_sixin_data "
               "WHERE session_id = %s ORDER BY date_num DESC, id DESC LIMIT 5")
        print(await client.query(sql, ("session_a",)))
    finally:
        await client.close()


if __name__ == "__main__":
    asyncio.run(main())