long_articles_agentic_supply/supply/infra/database/mysql.py

"""MySQL 异步后端 —— 基于 aiomysql 连接池"""

import logging
from typing import Any, Optional

from aiomysql import create_pool, DictCursor

from supply.infra.database.ports import SqlBackend

logger = logging.getLogger(__name__)


class MySQLBackend(SqlBackend):
    """MySQL 异步连接池后端"""

    def __init__(self, name: str, config: Any):
        super().__init__(name, config)
        self._pool = None

    # ==================== 生命周期 ====================

    async def init(self) -> None:
        if self._initialized:
            return
        self._pool = await create_pool(
            host=self.config.host,
            port=self.config.port,
            user=self.config.user,
            password=self.config.password,
            db=self.config.db,
            charset=getattr(self.config, "charset", "utf8mb4"),
            minsize=getattr(self.config, "minsize", 5),
            maxsize=getattr(self.config, "maxsize", 20),
            cursorclass=DictCursor,
            autocommit=True,
        )
        self._initialized = True
        logger.info("MySQL backend [%s] 连接池创建成功", self.name)

    async def close(self) -> None:
        if self._pool is None:
            return
        self._pool.close()
        await self._pool.wait_closed()
        self._pool = None
        self._initialized = False
        logger.info("MySQL backend [%s] 连接池已关闭", self.name)

    async def health(self) -> bool:
        try:
            await self.fetch_one("SELECT 1")
            return True
        except Exception:
            return False

    # ==================== SQL 操作 ====================

    async def fetch(
        self,
        query: str,
        params: Optional[tuple] = None,
    ) -> list[dict]:
        self._ensure_ready()
        async with self._pool.acquire() as conn:
            async with conn.cursor(DictCursor) as cursor:
                await cursor.execute(query, params)
                return await cursor.fetchall()

    async def fetch_one(
        self,
        query: str,
        params: Optional[tuple] = None,
    ) -> Optional[dict]:
        self._ensure_ready()
        async with self._pool.acquire() as conn:
            async with conn.cursor(DictCursor) as cursor:
                await cursor.execute(query, params)
                return await cursor.fetchone()

    async def execute(
        self,
        query: str,
        params: Any = None,
        *,
        batch: bool = False,
    ) -> int:
        self._ensure_ready()
        async with self._pool.acquire() as conn:
            async with conn.cursor() as cursor:
                try:
                    if batch:
                        await cursor.executemany(query, params)
                    else:
                        await cursor.execute(query, params)
                    affected = cursor.rowcount
                    await conn.commit()
                    return affected
                except Exception:
                    await conn.rollback()
                    raise

    # ==================== 内部 ====================

    def _ensure_ready(self):
        if not self._initialized or self._pool is None:
            raise RuntimeError(
                f"MySQL backend [{self.name}] 未初始化，请先调用 init()"
            )