Tolan/backend/database.py

"""
数据库连接管理模块
"""
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
    
    async def connect(self):
        """初始化数据库连接"""
        try:
            await self._create_mysql_pool()
            await self._create_redis_client()
        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:
            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 _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, "执行更新")
    
    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()