Maas-Model-platform/backend/app/core/async_database.py

"""
异步数据库配置模块

提供异步数据库连接、会话管理和依赖注入功能
使用 asyncpg 驱动实现高性能异步数据库操作
"""
import os
import time
import logging
from pathlib import Path
from typing import List, Dict, Any, Type, Optional
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy import text, inspect, event
from sqlalchemy.dialects.postgresql import insert
from dotenv import load_dotenv

# 配置慢查询日志记录器
logger = logging.getLogger(__name__)

# 加载环境变量（从backend目录的.env文件）
env_path = Path(__file__).parent.parent.parent / '.env'
load_dotenv(dotenv_path=env_path)

# 从环境变量读取数据库配置
DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_USER = os.getenv("DB_USER", "postgres")
DB_PASSWORD = os.getenv("DB_PASSWORD", "")
DB_NAME = os.getenv("DB_NAME", "model_square")
DB_POOL_SIZE = int(os.getenv("DB_POOL_SIZE", "20"))
DB_MAX_OVERFLOW = int(os.getenv("DB_MAX_OVERFLOW", "40"))
DB_POOL_TIMEOUT = int(os.getenv("DB_POOL_TIMEOUT", "30"))
DB_POOL_RECYCLE = int(os.getenv("DB_POOL_RECYCLE", "1800"))  # 30分钟回收连接

# 慢查询阈值配置（毫秒）
SLOW_QUERY_THRESHOLD_MS = int(os.getenv("SLOW_QUERY_THRESHOLD_MS", "1000"))

# 构建异步数据库连接URL（使用 asyncpg 驱动，对密码进行URL编码）
from urllib.parse import quote_plus
ASYNC_DATABASE_URL = f"postgresql+asyncpg://{DB_USER}:{quote_plus(DB_PASSWORD)}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

# 创建异步数据库引擎
async_engine = create_async_engine(
    ASYNC_DATABASE_URL,
    pool_size=DB_POOL_SIZE,
    max_overflow=DB_MAX_OVERFLOW,
    pool_timeout=DB_POOL_TIMEOUT,
    pool_recycle=DB_POOL_RECYCLE,
    pool_pre_ping=True,  # 连接前检查连接是否有效
    echo=False  # 生产环境关闭SQL日志
)


# 慢查询日志事件监听器
@event.listens_for(async_engine.sync_engine, "before_cursor_execute")
def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
    """
    在SQL执行前记录开始时间
    """
    conn.info.setdefault("query_start_time", []).append(time.time())


@event.listens_for(async_engine.sync_engine, "after_cursor_execute")
def after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
    """
    在SQL执行后检查执行时间，超过阈值则记录慢查询日志
    """
    start_time_list = conn.info.get("query_start_time", [])
    if start_time_list:
        total_time_ms = (time.time() - start_time_list.pop()) * 1000
        if total_time_ms > SLOW_QUERY_THRESHOLD_MS:
            # 截断SQL语句避免日志过长
            truncated_sql = statement[:500] + "..." if len(statement) > 500 else statement
            # 截断参数避免日志过长
            params_str = str(parameters)[:200] + "..." if len(str(parameters)) > 200 else str(parameters)
            logger.warning(
                f"慢查询 ({total_time_ms:.2f}ms): {truncated_sql} | 参数: {params_str}"
            )

# 创建异步会话工厂
AsyncSessionLocal = async_sessionmaker(
    bind=async_engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False
)


async def get_async_db():
    """
    异步数据库会话依赖注入函数
    
    用于FastAPI的依赖注入系统，自动管理异步数据库会话的生命周期
    确保请求结束后会话被正确关闭
    
    Yields:
        AsyncSession: 异步数据库会话对象
    """
    async with AsyncSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()


async def get_async_db_context():
    """
    异步数据库会话上下文管理器
    
    用于非依赖注入场景（如后台任务、定时任务等）
    
    Returns:
        AsyncSession: 异步数据库会话对象
    """
    return AsyncSessionLocal()


async def async_bulk_insert(
    session: AsyncSession,
    model: Type,
    data_list: List[Dict[str, Any]],
    batch_size: int = 1000
) -> int:
    """
    异步批量插入
    
    使用 PostgreSQL 的批量插入优化，减少数据库往返次数
    
    Args:
        session: 异步数据库会话
        model: SQLAlchemy 模型类
        data_list: 要插入的数据列表
        batch_size: 每批次插入的记录数，默认1000
    
    Returns:
        int: 插入的记录数
    """
    if not data_list:
        return 0
    
    total_inserted = 0
    
    # 分批处理，避免单次插入数据量过大
    for i in range(0, len(data_list), batch_size):
        batch = data_list[i:i + batch_size]
        stmt = insert(model).values(batch)
        await session.execute(stmt)
        total_inserted += len(batch)
    
    return total_inserted


async def async_bulk_update(
    session: AsyncSession,
    model: Type,
    data_list: List[Dict[str, Any]],
    key_field: str = "id",
    batch_size: int = 500
) -> int:
    """
    异步批量更新
    
    使用 PostgreSQL 的 ON CONFLICT 实现高效批量更新
    
    Args:
        session: 异步数据库会话
        model: SQLAlchemy 模型类
        data_list: 要更新的数据列表（必须包含 key_field）
        key_field: 用于匹配记录的字段名，默认为 "id"
        batch_size: 每批次更新的记录数，默认500
    
    Returns:
        int: 更新的记录数
    """
    if not data_list:
        return 0
    
    total_updated = 0
    
    # 获取模型的所有列名（排除主键用于更新）
    mapper = inspect(model)
    all_columns = [c.key for c in mapper.columns]
    
    # 分批处理
    for i in range(0, len(data_list), batch_size):
        batch = data_list[i:i + batch_size]
        
        # 构建 upsert 语句
        stmt = insert(model).values(batch)
        
        # 获取需要更新的列（排除 key_field）
        update_columns = {
            col: stmt.excluded[col] 
            for col in all_columns 
            if col != key_field and col in batch[0]
        }
        
        if update_columns:
            stmt = stmt.on_conflict_do_update(
                index_elements=[key_field],
                set_=update_columns
            )
            await session.execute(stmt)
            total_updated += len(batch)
    
    return total_updated


async def async_execute_raw(
    session: AsyncSession,
    sql: str,
    params: Optional[Dict[str, Any]] = None
) -> Any:
    """
    异步执行原生 SQL
    
    用于执行复杂查询或特定的数据库操作
    
    Args:
        session: 异步数据库会话
        sql: SQL 语句
        params: SQL 参数字典
    
    Returns:
        Any: 查询结果
    """
    stmt = text(sql)
    if params:
        result = await session.execute(stmt, params)
    else:
        result = await session.execute(stmt)
    return result


async def async_bulk_insert_returning(
    session: AsyncSession,
    model: Type,
    data_list: List[Dict[str, Any]],
    returning_fields: List[str] = None,
    batch_size: int = 1000
) -> List[Dict[str, Any]]:
    """
    异步批量插入并返回插入的记录
    
    Args:
        session: 异步数据库会话
        model: SQLAlchemy 模型类
        data_list: 要插入的数据列表
        returning_fields: 需要返回的字段列表，默认返回所有字段
        batch_size: 每批次插入的记录数
    
    Returns:
        List[Dict[str, Any]]: 插入的记录列表
    """
    if not data_list:
        return []
    
    results = []
    mapper = inspect(model)
    
    # 确定返回的字段
    if returning_fields:
        return_columns = [getattr(model, f) for f in returning_fields if hasattr(model, f)]
    else:
        return_columns = [c for c in mapper.columns]
    
    # 分批处理
    for i in range(0, len(data_list), batch_size):
        batch = data_list[i:i + batch_size]
        stmt = insert(model).values(batch).returning(*return_columns)
        result = await session.execute(stmt)
        rows = result.fetchall()
        
        # 转换为字典列表
        for row in rows:
            if returning_fields:
                results.append(dict(zip(returning_fields, row)))
            else:
                results.append(dict(zip([c.key for c in mapper.columns], row)))
    
    return results