""" 数据库连接管理模块 """ import asyncio import aiomysql import redis.asyncio as aioredis from typing import Optional, Callable, Awaitable from contextlib import asynccontextmanager from config import settings import logging logger = logging.getLogger(__name__) MYSQL_RETRY_EXCEPTIONS = ( aiomysql.OperationalError, aiomysql.InterfaceError, aiomysql.InternalError, ConnectionError, ConnectionResetError, ConnectionAbortedError, BrokenPipeError, TimeoutError, asyncio.TimeoutError, ) class DatabaseManager: """数据库连接管理器（MySQL + Redis）""" def __init__(self): self.mysql_pool: Optional[aiomysql.Pool] = None self.redis_client: Optional[aioredis.Redis] = None self._mysql_lock = asyncio.Lock() self._redis_lock = asyncio.Lock() self._last_mysql_healthcheck: float = 0.0 self._last_redis_healthcheck: float = 0.0 self._healthcheck_task: Optional[asyncio.Task] = None self._healthcheck_running = False async def connect(self): """初始化数据库连接""" try: await self._create_mysql_pool() await self._create_redis_client() # 🔥 启动后台健康检查任务 self._start_healthcheck_task() except Exception as e: logger.error(f"❌ 数据库连接失败: {e}") raise async def _create_mysql_pool(self): async with self._mysql_lock: if self.mysql_pool: self.mysql_pool.close() await self.mysql_pool.wait_closed() self.mysql_pool = None 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, connect_timeout=settings.MYSQL_CONNECT_TIMEOUT, # 注意：aiomysql 不支持 read_timeout 参数，已移除 pool_recycle=settings.MYSQL_POOL_RECYCLE, ) self._last_mysql_healthcheck = 0.0 logger.info( f"✅ MySQL连接池已创建: {settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATABASE}" ) async def _create_redis_client(self): async with self._redis_lock: if self.redis_client: await self.redis_client.close() self.redis_client = None redis_protocol = "rediss" if settings.REDIS_SSL else "redis" redis_url = ( f"{redis_protocol}://:{settings.REDIS_PASSWORD}@" f"{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, socket_timeout=30, # 🔥 增加超时时间，避免短时网络波动断开连接 retry_on_timeout=True, health_check_interval=30, socket_keepalive=True, # 🔥 启用 TCP keepalive，保持长连接（Redis 3.2+ 默认启用） ) await self.redis_client.ping() logger.info( f"✅ Redis连接成功: {settings.REDIS_HOST}:{settings.REDIS_PORT}/{settings.REDIS_DB}" ) async def disconnect(self): """关闭数据库连接""" try: # 🔥 停止健康检查任务 self._stop_healthcheck_task() if self.mysql_pool: self.mysql_pool.close() await self.mysql_pool.wait_closed() self.mysql_pool = None logger.info("✅ MySQL连接池已关闭") if self.redis_client: await self.redis_client.close() self.redis_client = None logger.info("✅ Redis连接已关闭") except Exception as e: logger.error(f"❌ 关闭数据库连接失败: {e}") async def _ensure_mysql_pool(self): """确保MySQL连接池可用（定期健康检查，保持长连接）""" if not self.mysql_pool: await self._create_mysql_pool() return loop = asyncio.get_running_loop() now = loop.time() # 🔥 健康检查间隔：只在超过间隔时检查，避免频繁检查 if now - self._last_mysql_healthcheck < settings.MYSQL_HEALTHCHECK_INTERVAL: return # 🔥 只做健康检查，不重建连接（保持长连接） try: async with self.mysql_pool.acquire() as conn: await conn.ping() self._last_mysql_healthcheck = now # 健康检查成功，连接正常，不重建 except Exception as exc: # 🔥 只有在健康检查失败时才重建连接（连接已断开） logger.warning(f"⚠️ MySQL健康检查失败，连接可能已断开，准备重建: {exc}") await self._create_mysql_pool() self._last_mysql_healthcheck = loop.time() async def _ensure_redis_client(self): """确保Redis连接可用（定期健康检查，自动重连）""" if not self.redis_client: await self._create_redis_client() return loop = asyncio.get_running_loop() now = loop.time() # 🔥 Redis健康检查间隔（从配置读取） if now - self._last_redis_healthcheck < settings.REDIS_HEALTHCHECK_INTERVAL: return # 🔥 执行健康检查 try: await self.redis_client.ping() self._last_redis_healthcheck = now # 健康检查成功，连接正常 except Exception as exc: # 🔥 健康检查失败，重建连接 logger.warning(f"⚠️ Redis健康检查失败，连接可能已断开，准备重建: {exc}") await self._create_redis_client() self._last_redis_healthcheck = loop.time() def _start_healthcheck_task(self): """启动后台健康检查任务""" if self._healthcheck_running: return self._healthcheck_running = True async def healthcheck_loop(): """定期健康检查循环""" while self._healthcheck_running: try: await asyncio.sleep(300) # 每5分钟检查一次 if not self._healthcheck_running: break # 检查MySQL await self._ensure_mysql_pool() # 检查Redis await self._ensure_redis_client() except Exception as e: logger.error(f"❌ 健康检查任务异常: {e}") self._healthcheck_task = asyncio.create_task(healthcheck_loop()) logger.info("✅ 数据库健康检查任务已启动") def _stop_healthcheck_task(self): """停止后台健康检查任务""" self._healthcheck_running = False if self._healthcheck_task and not self._healthcheck_task.done(): self._healthcheck_task.cancel() logger.info("✅ 数据库健康检查任务已停止") async def _run_with_retry(self, operation: Callable[[], Awaitable], description: str): """执行MySQL操作并在连接异常时重试一次""" for attempt in range(2): try: await self._ensure_mysql_pool() return await operation() except MYSQL_RETRY_EXCEPTIONS as exc: if attempt == 1: logger.error(f"❌ {description} 失败，已重试: {exc}") raise logger.warning(f"⚠️ {description} 连接异常，尝试重建连接: {exc}") await self._create_mysql_pool() @asynccontextmanager async def get_mysql_connection(self): """获取MySQL连接（上下文管理器）""" await self._ensure_mysql_pool() 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）""" params = params or {} async def _operation(): async with self.get_mysql_connection() as cursor: await cursor.execute(query, params) return await cursor.fetchall() return await self._run_with_retry(_operation, "执行查询") async def execute_one(self, query: str, params: dict = None): """执行查询并返回一条记录""" params = params or {} async def _operation(): async with self.get_mysql_connection() as cursor: await cursor.execute(query, params) return await cursor.fetchone() return await self._run_with_retry(_operation, "执行单条查询") async def execute_update(self, query: str, params: dict = None): """执行更新（INSERT/UPDATE/DELETE）""" params = params or {} async def _operation(): async with self.get_mysql_connection() as cursor: await cursor.execute(query, params) return cursor.lastrowid return await self._run_with_retry(_operation, "执行更新") async def get_redis(self) -> aioredis.Redis: """获取Redis客户端（自动健康检查）""" await self._ensure_redis_client() 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()