Mysql接口开发-aiomysql

Norles

已于 2024-08-14 13:40:11 修改

阅读量847

点赞数 6

文章标签： mysql 数据库

于 2024-08-14 13:32:40 首次发布

本文链接：https://blog.csdn.net/Norles/article/details/141188540

版权

概述

aiomysql 是一个用于 Python 的异步 MySQL 客户端库，它允许你使用 asyncio 编写异步代码来与 MySQL 数据库进行交互

这段代码定义了一个名为 `ClientManager` 的类，用于管理数据库连接和执行SQL查询。它使用了 `aiomysql` 库来与MySQL数据库进行异步交互
1. **导入模块**:
- `reprlib` 用于生成对象的字符串表示。
- `asyncio` 用于异步编程。
- `loguru` 用于日志记录。
- `cachetools` 用于缓存管理。
- `pandas` 用于数据分析和数据处理。
- `aiomysql` 用于异步MySQL数据库交互。
- `pymysql.constants` 用于定义MySQL常量。
- `deepfos` 和 `deepfos.api.datatable` 可能是一个自定义的框架或模块。
2. **全局设置**:
- `aRepr.maxstring = 200`：设置`reprlib`的最大字符串长度。
3. **`Client` 类**:
- 管理MySQL连接池。
- 使用`asyncio.Event`来等待连接池初始化完成。
- 包含连接池、状态、初始化事件、租约管理器和密码。
- `__init__`方法初始化属性。
- `init_pool`方法初始化连接池。
- `ensure_connected`方法确保连接池已建立。
- `execute`、`select`、`query_dataframe`和`trxn_execute`方法执行不同的SQL操作。
- `close`方法关闭连接池。
4. **`ClientManager` 类**:
- 使用`TTLCache`管理`Client`实例的缓存。
- `new_client`方法获取一个新的或缓存的`Client`实例。
5. **`LeaseManager` 类**:
- 继承自`AbsLeaseManager`。
- 管理数据库凭据的租约。
- `new_api`方法创建一个新的MySQLAPI实例。

这段代码定义了一个名为 `AbsLeaseManager` 的抽象类，用于管理数据库连接的租约
**`AbsLeaseManager` 类**:
- `DB` 类变量，用于存储数据库类型。
- `__init__` 方法初始化租约管理器，包括设置最后续租时间、续租间隔、任务标识、租约信息和续租函数。
- `new_api` 方法异步地创建一个新的API实例。
- `_try_renewal` 方法尝试续租，如果失败，会重试一次。
- `renew` 方法执行续租操作，如果失败，会重试。
- `loop` 方法启动一个无限循环，不断尝试续租。
- `schedule` 方法安排续租任务。
- `cancel` 方法取消续租任务。
这个 `AbsLeaseManager` 类的主要作用是管理数据库连接的租约，通过异步方式进行续租操作，确保数据库连接的有效性。它使用重试机制来处理续租过程中的异常情况。

6. **辅助函数**:
- `select`、`execute`、`query_dataframe`和`trxn_execute`函数使用`ClientManager`获取`Client`实例并执行相应的数据库操作。
7. **日志记录和异常处理**:
- 使用`logger`进行日志记录。
- 定义`DBConnecetionError`作为自定义的连接错误。
这个模块提供了对MySQL数据库的异步访问，并且通过连接池和缓存机制来提高性能和效率。它还包含了一个事务执行的方法，可以一次执行多个SQL语句。

from reprlib import aRepr
import asyncio

from loguru import logger
from cachetools import TTLCache
import pandas as pd
import aiomysql
from pymysql.constants import FIELD_TYPE

from deepfos import OPTION
from deepfos.api.datatable import MySQLAPI
from deepfos.lib.constant import UNSET
from deepfos.lib.asynchronous import register_on_loop_shutdown
from .utils import (
    decrypt, AbsLeaseManager, ACCOUNT_EXPIRE,
    PENDING, INITIALIZED, INITIALIZING, DBConnecetionError
)

aRepr.maxstring = 200


# cdef class Client:
#     cdef:
#         object  pool
#         int status
#         object  inited
#         object  lease
#         bytes secret
class Client:
    def __init__(self):
        self.pool = UNSET
        self.status = PENDING
        self.inited = asyncio.Event()
        # 租约
        self.lease = LeaseManager(ACCOUNT_EXPIRE / 3)
        self.secret = "!ABCD-EFGH-IJKL@".encode()
        # 注册一个回调函数
        # self.close，当事件循环关闭时执行。这个函数可能用于关闭数据库连接池或执行其他清理工作。
        register_on_loop_shutdown(self.close, True)

    async def init_pool(self):
        if self.status == INITIALIZED:
            return self.pool

        if self.status == PENDING:
            # 如果 # self.status # 等于# PENDING，则将其更新为
            # INITIALIZING，这可能是一个状态机，用于跟踪连接池的初始化过程。
            self.status = INITIALIZING

        #     如果 self.status 不等于 INITIALIZED 且不等于 PENDING，则说明连接池正在初始化中。
        else:
            await self.inited.wait()
            if self.status != INITIALIZED:
                raise RuntimeError("Failed to initialze connection pool.")
            return self.pool

        # 使用 self.lease 对象的 renew 方法获取新的数据库配置信息。self.lease 是一个 LeaseManager 实例，用于管理租约。
        conf = await self.lease.renew()
        self.lease.schedule(slow_start=True)

        try:
            # 使用 aiomysql.create_pool 异步创建一个连接池。
            # 参数包括最小连接数 minsize、最大连接数 maxsize、主机 host、端口 port、用户名 user、密码 password（需要解密）和数据库名 db。
            self.pool = await aiomysql.create_pool(
                minsize=5,
                maxsize=10,
                host=conf.host,
                port=conf.port,
                user=conf.name,
                password=decrypt(self.secret, conf.password),
                db=conf.dbName,
            )
        except Exception as e:
            self.status = PENDING
            raise DBConnecetionError(e) from None
        else:
            self.status = INITIALIZED
            self.inited.set()
            return self.pool

    async def ensure_connected(self):
        max_retry = 3
        retries = 0
        interval = 0.5

        while self.pool is UNSET or self.pool._closed:  # noqa
            retries += 1
            try:
                await self.init_pool()
            except DBConnecetionError:
                if retries > max_retry:
                    self.inited.set()
                    raise
                logger.exception(f'Failed to get connection pool, '
                                 f'starting {retries} times retry.')
                await asyncio.sleep(interval)
                interval *= 2
            except Exception:
                self.inited.set()
                self.status = PENDING
                raise

        return self.pool

    async def execute(self, sql):
        pool = await self.ensure_connected()
        async with pool.acquire() as conn:
            async with conn.cursor() as cur:
                logger.opt(lazy=True).debug("Run sql: {sql}",
                                            sql=lambda: aRepr.repr(sql))
                await cur.execute(sql)
            await conn.commit()
            return conn.affected_rows()

    async def select(self, sql):
        pool = await self.ensure_connected()
        async with pool.acquire() as conn:
            async with conn.cursor() as cur:
                logger.opt(lazy=True).debug("Run sql: {sql}",
                                            sql=lambda: aRepr.repr(sql))
                await cur.execute(sql)
                return await cur.fetchall()

    async def query_dataframe(self, sql):
        pool = await self.ensure_connected()
        async with pool.acquire() as conn:
            async with conn.cursor() as cur:
                logger.opt(lazy=True).debug("Run sql: {sql}",
                                            sql=lambda: aRepr.repr(sql))
                await cur.execute(sql)
                columns = [d[0] for d in cur.description]
                decimal_cols = [
                    d[0] for d in cur.description
                    if d[1] == FIELD_TYPE.NEWDECIMAL
                ]
                data = await cur.fetchall()
                df = pd.DataFrame(data=data, columns=columns)
                for col in decimal_cols:
                    df[col] = pd.to_numeric(df[col], downcast='float')
                return df

    async def trxn_execute(self, sqls):
        pool = await self.ensure_connected()
        async with pool.acquire() as conn:
            async with conn.cursor() as cur:
                logger.opt(lazy=True).debug("Run sql: {sql}",
                                            sql=lambda: aRepr.repr(sqls))
                for sql in sqls:
                    await cur.execute(sql)
            await conn.commit()
            return conn.affected_rows()

    async def close(self):
        self.lease.cancel()
        self.status = PENDING

        if self.pool is not UNSET:
            self.pool.close()
            await self.pool.wait_closed()
            self.pool = UNSET


# cdef class ClientManager:
class ClientManager:
    cache = TTLCache(maxsize=10, ttl=1800)

    def new_client(self):
        key = (
            OPTION.api.header.get('space'),
            OPTION.api.header.get('app'),
        )

        if key not in self.cache:
            self.cache[key] = Client()

        return self.cache[key]


class LeaseManager(AbsLeaseManager):
    DB = 'mysql'

    async def new_api(self):
        return await MySQLAPI(version=1.0, sync=False)


async def select(sql):
    return await ClientManager().new_client().select(sql)


async def execute(sql):
    return await ClientManager().new_client().execute(sql)


async def query_dataframe(sql):
    return await ClientManager().new_client().query_dataframe(sql)


async def trxn_execute(sqls):
    return await ClientManager().new_client().trxn_execute(sqls)