""" 数据库会话管理和依赖注入。 """ from contextlib import contextmanager from typing import Generator from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker, Session from sqlalchemy.pool import QueuePool from app.config import Settings class DatabaseManager: """ 数据库连接和会话管理器。 """ def __init__(self, settings: Settings): self.settings = settings self.engine = None self.SessionLocal = None self._initialize_engine() def _initialize_engine(self) -> None: """ 使用连接池初始化数据库引擎。 """ self.engine = create_engine( self.settings.database_url, poolclass=QueuePool, pool_size=5, max_overflow=10, pool_pre_ping=True, # 使用前验证连接 pool_recycle=300, # 5 分钟后回收连接(之前是 3600) pool_reset_on_return="rollback", # 返回时重置连接 pool_timeout=30, # 等待连接最多 30 秒 connect_args={ "connect_timeout": 10, "read_timeout": 30, "write_timeout": 30, "charset": "utf8mb4", "autocommit": False, }, echo=self.settings.debug, isolation_level="READ COMMITTED", ) self.SessionLocal = sessionmaker( autocommit=False, autoflush=False, bind=self.engine, ) def get_db(self) -> Generator[Session, None, None]: """ 用于数据库会话的 FastAPI 依赖项。 Yields: Session: SQLAlchemy 数据库会话 """ db = self.SessionLocal() try: yield db except Exception: db.rollback() raise finally: try: db.close() except Exception: # 忽略关闭时的错误(连接可能已经关闭) pass @contextmanager def session_scope(self) -> Generator[Session, None, None]: """ 数据库会话的上下文管理器。 Yields: Session: SQLAlchemy 数据库会话 Example: with db_manager.session_scope() as session: session.query(Model).all() """ session = self.SessionLocal() try: yield session session.commit() except Exception: session.rollback() raise finally: session.close() def create_all_tables(self) -> None: """ 创建所有数据库表(用于测试/开发)。 """ from app.db.models import Base Base.metadata.create_all(bind=self.engine) def drop_all_tables(self) -> None: """ 删除所有数据库表(用于测试)。 """ from app.db.models import Base Base.metadata.drop_all(bind=self.engine) def close(self) -> None: """ 关闭数据库引擎。 """ if self.engine: try: self.engine.dispose() except Exception: # 忽略处理时的错误 pass # 全局数据库管理器实例(在 main.py 中初始化) db_manager: DatabaseManager = None def get_db_manager() -> DatabaseManager: """ 获取全局数据库管理器实例。 Returns: DatabaseManager: 数据库管理器实例 """ return db_manager def init_db_manager(settings: Settings) -> DatabaseManager: """ 初始化全局数据库管理器。 Args: settings: 应用程序设置 Returns: DatabaseManager: 初始化的数据库管理器 """ global db_manager db_manager = DatabaseManager(settings) return db_manager def get_db() -> Generator[Session, None, None]: """ 用于 FastAPI 依赖注入的全局 get_db 函数。 Yields: Session: SQLAlchemy 数据库会话 """ if db_manager is None: raise RuntimeError("Database manager not initialized. Call init_db_manager() first.") yield from db_manager.get_db()