Tolan/backend/database.py

"""
数据库连接管理模块
"""
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()