"""数据库连接和会话管理 - PostgreSQL 多用户数据隔离""" import asyncio from typing import Dict, Any from datetime import datetime from sqlalchemy import select, text from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy.orm import declarative_base from fastapi import Request, HTTPException from app.config import settings from app.logger import get_logger logger = get_logger(__name__) # 创建基类 Base = declarative_base() # 导入所有模型,确保 Base.metadata 能够发现它们 # 这必须在 Base 创建之后、init_db 之前导入 from app.models import ( Project, Outline, Character, Chapter, GenerationHistory, Settings, WritingStyle, ProjectDefaultStyle, RelationshipType, CharacterRelationship, Organization, OrganizationMember, StoryMemory, PlotAnalysis, AnalysisTask, BatchGenerationTask, RegenerationTask, Career, CharacterCareer, User, MCPPlugin, PromptTemplate ) # 引擎缓存:每个用户一个引擎 _engine_cache: Dict[str, Any] = {} # 锁管理:用于保护引擎创建过程 _engine_locks: Dict[str, asyncio.Lock] = {} _cache_lock = asyncio.Lock() # 会话统计(用于监控连接泄漏) _session_stats = { "created": 0, "closed": 0, "active": 0, "errors": 0, "generator_exits": 0, "last_check": None } async def get_engine(user_id: str): """获取或创建用户专属的数据库引擎(线程安全) PostgreSQL: 所有用户共享一个数据库,通过user_id字段隔离数据 Args: user_id: 用户ID Returns: 用户专属的异步引擎 """ # PostgreSQL模式:所有用户共享同一个引擎 cache_key = "shared_postgres" if cache_key in _engine_cache: return _engine_cache[cache_key] async with _cache_lock: if cache_key not in _engine_cache: # 检测数据库类型 is_sqlite = 'sqlite' in settings.database_url.lower() # 基础引擎参数 engine_args = { "echo": settings.database_echo_pool, "echo_pool": settings.database_echo_pool, "future": True, } if is_sqlite: # SQLite 配置(使用 NullPool,不支持连接池参数) engine_args["connect_args"] = { "check_same_thread": False, "timeout": 30.0, # 等待锁释放的超时时间(秒) } # 启用连接前检测以支持更好的并发 engine_args["pool_pre_ping"] = True logger.info("📊 使用 SQLite 数据库(NullPool,超时30秒,WAL模式)") else: # PostgreSQL 配置(完整连接池支持) connect_args = { "server_settings": { "application_name": settings.app_name, "jit": "off", "search_path": "public", }, "command_timeout": 60, "statement_cache_size": 500, } engine_args.update({ "pool_size": settings.database_pool_size, "max_overflow": settings.database_max_overflow, "pool_timeout": settings.database_pool_timeout, "pool_pre_ping": settings.database_pool_pre_ping, "pool_recycle": settings.database_pool_recycle, "pool_use_lifo": settings.database_pool_use_lifo, "pool_reset_on_return": settings.database_pool_reset_on_return, "max_identifier_length": settings.database_max_identifier_length, "connect_args": connect_args }) total_connections = settings.database_pool_size + settings.database_max_overflow estimated_concurrent_users = total_connections * 2 logger.info( f"📊 PostgreSQL 连接池配置:\n" f" ├─ 核心连接: {settings.database_pool_size}\n" f" ├─ 溢出连接: {settings.database_max_overflow}\n" f" ├─ 总连接数: {total_connections}\n" f" ├─ 获取超时: {settings.database_pool_timeout}秒\n" f" ├─ 连接回收: {settings.database_pool_recycle}秒\n" f" └─ 预估并发: {estimated_concurrent_users}+用户" ) engine = create_async_engine(settings.database_url, **engine_args) _engine_cache[cache_key] = engine # 如果是 SQLite,启用 WAL 模式以支持读写并发 if is_sqlite: try: from sqlalchemy import event from sqlalchemy.pool import NullPool @event.listens_for(engine.sync_engine, "connect") def set_sqlite_pragma(dbapi_conn, connection_record): cursor = dbapi_conn.cursor() cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA synchronous=NORMAL") cursor.execute("PRAGMA cache_size=-64000") # 64MB 缓存 cursor.execute("PRAGMA busy_timeout=30000") # 30秒超时 cursor.close() logger.info("✅ SQLite WAL 模式已启用(支持读写并发)") except Exception as e: logger.warning(f"⚠️ 启用 WAL 模式失败: {e},使用默认配置") return _engine_cache[cache_key] async def get_db(request: Request): """获取数据库会话的依赖函数 从 request.state.user_id 获取用户ID,然后返回该用户的数据库会话 """ user_id = getattr(request.state, "user_id", None) if not user_id: raise HTTPException(status_code=401, detail="未登录或用户ID缺失") engine = await get_engine(user_id) AsyncSessionLocal = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False ) session = AsyncSessionLocal() session_id = id(session) global _session_stats _session_stats["created"] += 1 _session_stats["active"] += 1 # logger.debug(f"📊 会话创建 [User:{user_id}][ID:{session_id}] - 活跃:{_session_stats['active']}, 总创建:{_session_stats['created']}, 总关闭:{_session_stats['closed']}") try: yield session if session.in_transaction(): await session.rollback() except GeneratorExit: _session_stats["generator_exits"] += 1 logger.warning(f"⚠️ GeneratorExit [User:{user_id}][ID:{session_id}] - SSE连接断开(总计:{_session_stats['generator_exits']}次)") try: if session.in_transaction(): await session.rollback() logger.info(f"✅ 事务已回滚 [User:{user_id}][ID:{session_id}](GeneratorExit)") except Exception as rollback_error: _session_stats["errors"] += 1 logger.error(f"❌ GeneratorExit回滚失败 [User:{user_id}][ID:{session_id}]: {str(rollback_error)}") except Exception as e: _session_stats["errors"] += 1 logger.error(f"❌ 会话异常 [User:{user_id}][ID:{session_id}]: {str(e)}") try: if session.in_transaction(): await session.rollback() logger.info(f"✅ 事务已回滚 [User:{user_id}][ID:{session_id}](异常)") except Exception as rollback_error: logger.error(f"❌ 异常回滚失败 [User:{user_id}][ID:{session_id}]: {str(rollback_error)}") raise finally: try: if session.in_transaction(): await session.rollback() logger.warning(f"⚠️ finally中发现未提交事务 [User:{user_id}][ID:{session_id}],已回滚") await session.close() _session_stats["closed"] += 1 _session_stats["active"] -= 1 _session_stats["last_check"] = datetime.now().isoformat() # logger.debug(f"📊 会话关闭 [User:{user_id}][ID:{session_id}] - 活跃:{_session_stats['active']}, 总创建:{_session_stats['created']}, 总关闭:{_session_stats['closed']}, 错误:{_session_stats['errors']}") # 使用优化后的会话监控阈值 if _session_stats["active"] > settings.database_session_leak_threshold: logger.error(f"🚨 严重告警:活跃会话数 {_session_stats['active']} 超过泄漏阈值 {settings.database_session_leak_threshold}!") elif _session_stats["active"] > settings.database_session_max_active: logger.warning(f"⚠️ 警告:活跃会话数 {_session_stats['active']} 超过警告阈值 {settings.database_session_max_active},可能存在连接泄漏!") elif _session_stats["active"] < 0: logger.error(f"🚨 活跃会话数异常: {_session_stats['active']},统计可能不准确!") except Exception as e: _session_stats["errors"] += 1 logger.error(f"❌ 关闭会话时出错 [User:{user_id}][ID:{session_id}]: {str(e)}", exc_info=True) try: await session.close() except: pass async def init_db(user_id: str = None): """ 初始化数据库(已弃用) ⚠️ 此函数已弃用,仅保留用于向后兼容 新的最佳实践: - 表结构管理: 使用 'alembic upgrade head' - 用户配置: Settings 在首次访问时自动创建(延迟初始化) Args: user_id: 用户ID (已不再使用) """ logger.warning( "⚠️ init_db() 已弃用且无实际作用!\n" " - 表结构: 由 Alembic 管理\n" " - 用户配置: Settings API 自动创建\n" " 建议移除此调用" ) async def close_db(): """关闭所有数据库连接""" try: logger.info("正在关闭所有数据库连接...") for user_id, engine in _engine_cache.items(): await engine.dispose() logger.info(f"用户 {user_id} 的数据库连接已关闭") _engine_cache.clear() logger.info("所有数据库连接已关闭") except Exception as e: logger.error(f"关闭数据库连接失败: {str(e)}", exc_info=True) raise async def get_database_stats(): """获取数据库连接和会话统计信息 Returns: dict: 包含数据库统计信息的字典 """ from app.config import settings # 获取连接池详细状态 pool_stats = {} cache_key = "shared_postgres" if cache_key in _engine_cache: engine = _engine_cache[cache_key] try: pool = engine.pool pool_stats = { "size": pool.size(), # 当前连接池大小 "checked_in": pool.checkedin(), # 可用连接数 "checked_out": pool.checkedout(), # 正在使用的连接数 "overflow": pool.overflow(), # 溢出连接数 "usage_percent": (pool.checkedout() / (settings.database_pool_size + settings.database_max_overflow)) * 100, } except Exception as e: logger.warning(f"获取连接池状态失败: {e}") pool_stats = {"error": str(e)} stats = { "session_stats": { "created": _session_stats["created"], "closed": _session_stats["closed"], "active": _session_stats["active"], "errors": _session_stats["errors"], "generator_exits": _session_stats["generator_exits"], "last_check": _session_stats["last_check"], }, "pool_stats": pool_stats, # 新增:连接池实时状态 "engine_cache": { "total_engines": len(_engine_cache), "engine_keys": list(_engine_cache.keys()), }, "config": { "database_type": "PostgreSQL", "pool_size": settings.database_pool_size, "max_overflow": settings.database_max_overflow, "total_connections": settings.database_pool_size + settings.database_max_overflow, "pool_timeout": settings.database_pool_timeout, "pool_recycle": settings.database_pool_recycle, "session_max_active_threshold": settings.database_session_max_active, "session_leak_threshold": settings.database_session_leak_threshold, }, "health": { "status": "healthy", "warnings": [], "errors": [], } } # 健康检查 if _session_stats["active"] > settings.database_session_leak_threshold: stats["health"]["status"] = "critical" stats["health"]["errors"].append( f"活跃会话数 {_session_stats['active']} 超过泄漏阈值 {settings.database_session_leak_threshold}" ) elif _session_stats["active"] > settings.database_session_max_active: stats["health"]["status"] = "warning" stats["health"]["warnings"].append( f"活跃会话数 {_session_stats['active']} 超过警告阈值 {settings.database_session_max_active}" ) if _session_stats["active"] < 0: stats["health"]["status"] = "error" stats["health"]["errors"].append(f"活跃会话数异常: {_session_stats['active']}") # 连接池使用率检查 if pool_stats and "usage_percent" in pool_stats: usage = pool_stats["usage_percent"] if usage > 90: stats["health"]["status"] = "warning" stats["health"]["warnings"].append(f"连接池使用率过高: {usage:.1f}%") elif usage > 95: stats["health"]["status"] = "critical" stats["health"]["errors"].append(f"连接池几乎耗尽: {usage:.1f}%") error_rate = (_session_stats["errors"] / max(_session_stats["created"], 1)) * 100 if error_rate > 5: if stats["health"]["status"] == "healthy": stats["health"]["status"] = "warning" stats["health"]["warnings"].append(f"会话错误率过高: {error_rate:.2f}%") stats["health"]["error_rate"] = f"{error_rate:.2f}%" return stats async def check_database_health(user_id: str = None) -> dict: """检查数据库连接健康状态 Args: user_id: 可选的用户ID,如果提供则检查特定用户的数据库 Returns: dict: 健康检查结果 """ result = { "healthy": True, "checks": {}, "timestamp": datetime.now().isoformat() } try: # 检查引擎是否存在 cache_key = "shared_postgres" if user_id: engine = await get_engine(user_id) else: if cache_key not in _engine_cache: result["checks"]["engine"] = {"status": "not_initialized", "healthy": True} return result engine = _engine_cache[cache_key] # 测试数据库连接 AsyncSessionLocal = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False ) async with AsyncSessionLocal() as session: # 执行简单查询测试连接 await session.execute(text("SELECT 1")) result["checks"]["connection"] = {"status": "ok", "healthy": True} # 检查连接池状态(仅PostgreSQL) if hasattr(engine.pool, 'size'): pool_status = { "size": engine.pool.size(), "checked_in": engine.pool.checkedin(), "checked_out": engine.pool.checkedout(), "overflow": engine.pool.overflow(), "healthy": True } # 连接池健康检查 if engine.pool.overflow() >= settings.database_max_overflow: pool_status["healthy"] = False pool_status["warning"] = "连接池溢出已满" result["healthy"] = False result["checks"]["pool"] = pool_status except Exception as e: result["healthy"] = False result["checks"]["error"] = { "status": "error", "message": str(e), "healthy": False } logger.error(f"数据库健康检查失败: {str(e)}", exc_info=True) return result async def reset_session_stats(): """重置会话统计信息(用于测试或维护)""" global _session_stats _session_stats = { "created": 0, "closed": 0, "active": 0, "errors": 0, "generator_exits": 0, "last_check": datetime.now().isoformat() } logger.info("✅ 会话统计信息已重置") return _session_stats