""" 数据库连接管理模块（SQLAlchemy 版本） """ import asyncio import time import traceback from typing import Optional from urllib.parse import quote_plus import redis.asyncio as aioredis from sqlalchemy.ext.asyncio import ( create_async_engine, AsyncSession, async_sessionmaker, AsyncEngine ) from sqlalchemy.orm import declarative_base from config import settings import logging logger = logging.getLogger(__name__) class DatabaseManager: """数据库连接管理器（MySQL + Redis）""" def __init__(self): self.engine: Optional[AsyncEngine] = None self.async_session_maker: Optional[async_sessionmaker] = None self.redis_client: Optional[aioredis.Redis] = None self._redis_lock = asyncio.Lock() 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_engine() await self._create_redis_client() # 🔥 启动后台健康检查任务 self._start_healthcheck_task() except Exception as e: logger.error(f"❌ 数据库连接失败: {e}") raise async def _create_mysql_engine(self): """创建 SQLAlchemy 异步引擎""" # 构建 MySQL 连接 URL encoded_user = quote_plus(settings.MYSQL_USER) encoded_password = quote_plus(settings.MYSQL_PASSWORD) mysql_url = ( f"mysql+aiomysql://{encoded_user}:{encoded_password}" f"@{settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATABASE}" f"?charset=utf8mb4" ) # 创建异步引擎（自动连接池 + 健康检查 + 超时配置） self.engine = create_async_engine( mysql_url, pool_size=settings.MYSQL_POOL_SIZE, max_overflow=settings.MYSQL_MAX_OVERFLOW, pool_pre_ping=True, # 🔥 自动健康检查（类似 ping） pool_recycle=settings.MYSQL_POOL_RECYCLE, # 连接回收 pool_timeout=5.0, # 🔥 连接池获取连接超时（秒），避免长时间等待 echo=False, # 生产环境关闭 SQL 日志 future=True, connect_args={ "connect_timeout": settings.MYSQL_CONNECT_TIMEOUT, # 🔥 连接超时（aiomysql 支持） # ⚠️ read_timeout 已移除：aiomysql 不支持此参数，读取超时由 pool_timeout 控制 }, ) # 创建异步 Session 工厂 self.async_session_maker = async_sessionmaker( self.engine, class_=AsyncSession, expire_on_commit=False, autoflush=False, autocommit=False, ) logger.info( f"✅ MySQL连接池已创建: {settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATABASE}" ) async def _create_redis_client(self): """创建 Redis 客户端""" 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, ) 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.engine: await self.engine.dispose() self.engine = None self.async_session_maker = 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}") def get_session(self) -> AsyncSession: """ 获取异步 Session（上下文管理器） 使用方式： async with db_manager.get_session() as session: user = await session.get(User, user_id) """ if not self.async_session_maker: raise RuntimeError("MySQL连接池未初始化") return self.async_session_maker() async def _ensure_redis_client(self): """确保Redis连接可用（定期健康检查，自动重连）""" start_time = time.time() if not self.redis_client: logger.info("🔧 Redis客户端不存在，创建新连接...") await self._create_redis_client() elapsed = (time.time() - start_time) * 1000 logger.info(f"✅ Redis客户端创建完成，耗时: {elapsed:.1f}ms") return loop = asyncio.get_running_loop() now = loop.time() time_since_last_check = now - self._last_redis_healthcheck # 🔥 Redis健康检查间隔（从配置读取） if time_since_last_check < settings.REDIS_HEALTHCHECK_INTERVAL: logger.debug(f"⏭️ Redis健康检查跳过（距离上次检查: {time_since_last_check:.1f}s < {settings.REDIS_HEALTHCHECK_INTERVAL}s）") return logger.info(f"🏥 开始Redis健康检查（距离上次检查: {time_since_last_check:.1f}s）...") # 🔥 执行健康检查（使用超时，避免长时间阻塞） try: # 使用asyncio.wait_for设置超时，避免健康检查阻塞太久 ping_start = time.time() await asyncio.wait_for(self.redis_client.ping(), timeout=1.0) ping_elapsed = (time.time() - ping_start) * 1000 self._last_redis_healthcheck = now elapsed = (time.time() - start_time) * 1000 logger.info(f"✅ Redis健康检查成功，PING耗时: {ping_elapsed:.1f}ms，总耗时: {elapsed:.1f}ms") except asyncio.TimeoutError: elapsed = (time.time() - start_time) * 1000 # 健康检查超时，可能连接有问题，但不在关键路径上阻塞 logger.warning(f"⚠️ Redis健康检查超时（耗时: {elapsed:.1f}ms），将在后台重建连接") # 🔥 在后台重建连接，不阻塞当前请求 asyncio.create_task(self._recreate_redis_client_background()) self._last_redis_healthcheck = loop.time() except Exception as exc: elapsed = (time.time() - start_time) * 1000 # 🔥 健康检查失败，在后台重建连接，不阻塞当前请求 logger.warning(f"⚠️ Redis健康检查失败（耗时: {elapsed:.1f}ms），将在后台重建连接: {exc}") asyncio.create_task(self._recreate_redis_client_background()) self._last_redis_healthcheck = loop.time() async def _recreate_redis_client_background(self): """后台重建Redis连接（不阻塞主流程）""" start_time = time.time() logger.info("🔧 开始后台重建Redis连接...") try: old_client = self.redis_client self.redis_client = None # 先清空，避免其他请求使用旧连接 await self._create_redis_client() if old_client: try: await old_client.close() except: pass elapsed = (time.time() - start_time) * 1000 logger.info(f"✅ Redis连接已在后台重建完成，耗时: {elapsed:.1f}ms") except Exception as e: elapsed = (time.time() - start_time) * 1000 logger.error(f"❌ 后台重建Redis连接失败，耗时: {elapsed:.1f}ms, 错误: {e}") logger.error(traceback.format_exc()) 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 # 检查Redis（MySQL由SQLAlchemy的pool_pre_ping自动处理） 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 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()