""" 数据库连接管理模块 """ import aiomysql import redis.asyncio as aioredis from typing import Optional from contextlib import asynccontextmanager from config import settings import logging logger = logging.getLogger(__name__) class DatabaseManager: """数据库连接管理器(MySQL + Redis)""" def __init__(self): self.mysql_pool: Optional[aiomysql.Pool] = None self.redis_client: Optional[aioredis.Redis] = None async def connect(self): """初始化数据库连接""" try: # 连接MySQL self.mysql_pool = await aiomysql.create_pool( host=settings.MYSQL_HOST, port=settings.MYSQL_PORT, user=settings.MYSQL_USER, password=settings.MYSQL_PASSWORD, db=settings.MYSQL_DATABASE, charset='utf8mb4', minsize=5, maxsize=settings.MYSQL_POOL_SIZE, autocommit=True, echo=False ) logger.info(f"✅ MySQL连接池已创建: {settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATABASE}") # 连接Redis # 根据SSL配置选择协议(rediss:// 或 redis://) redis_protocol = "rediss" if settings.REDIS_SSL else "redis" redis_url = f"{redis_protocol}://:{settings.REDIS_PASSWORD}@{settings.REDIS_HOST}:{settings.REDIS_PORT}/{settings.REDIS_DB}" self.redis_client = await aioredis.from_url( redis_url, encoding="utf-8", decode_responses=True, max_connections=settings.REDIS_MAX_CONNECTIONS, socket_connect_timeout=10, # 增加连接超时到10秒 socket_timeout=10, # 增加操作超时到10秒 retry_on_timeout=True, # 超时自动重试 health_check_interval=30 # 健康检查间隔30秒 ) # 测试Redis连接 await self.redis_client.ping() logger.info(f"✅ Redis连接成功: {settings.REDIS_HOST}:{settings.REDIS_PORT}/{settings.REDIS_DB}") except Exception as e: logger.error(f"❌ 数据库连接失败: {e}") raise async def disconnect(self): """关闭数据库连接""" try: if self.mysql_pool: self.mysql_pool.close() await self.mysql_pool.wait_closed() logger.info("✅ MySQL连接池已关闭") if self.redis_client: await self.redis_client.close() logger.info("✅ Redis连接已关闭") except Exception as e: logger.error(f"❌ 关闭数据库连接失败: {e}") @asynccontextmanager async def get_mysql_connection(self): """获取MySQL连接(上下文管理器)""" if not self.mysql_pool: raise RuntimeError("MySQL连接池未初始化") async with self.mysql_pool.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cursor: yield cursor async def execute_query(self, query: str, params: dict = None): """执行查询(SELECT)""" async with self.get_mysql_connection() as cursor: await cursor.execute(query, params or {}) return await cursor.fetchall() async def execute_one(self, query: str, params: dict = None): """执行查询并返回一条记录""" async with self.get_mysql_connection() as cursor: await cursor.execute(query, params or {}) return await cursor.fetchone() async def execute_update(self, query: str, params: dict = None): """执行更新(INSERT/UPDATE/DELETE)""" async with self.get_mysql_connection() as cursor: await cursor.execute(query, params or {}) return cursor.lastrowid def get_redis(self) -> aioredis.Redis: """获取Redis客户端""" if not self.redis_client: raise RuntimeError("Redis客户端未初始化") return self.redis_client # 全局数据库管理器实例 db_manager = DatabaseManager() async def init_database(): """初始化数据库连接""" await db_manager.connect() async def close_database(): """关闭数据库连接""" await db_manager.disconnect()