Tolan/backend/database.py

"""
数据库连接管理模块（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,  # 连接回收
            echo=False,  # 生产环境关闭 SQL 日志
            future=True,
        )
        
        # 创建异步 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()