"""SQLAlchemy async engine + session。""" from __future__ import annotations from collections.abc import AsyncIterator from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase from src.core.config import settings class Base(DeclarativeBase): pass engine = create_async_engine( settings.DATABASE_URL, echo=False, pool_pre_ping=True, pool_size=settings.DB_POOL_SIZE, max_overflow=settings.DB_MAX_OVERFLOW, pool_timeout=settings.DB_POOL_TIMEOUT, pool_recycle=settings.DB_POOL_RECYCLE, ) AsyncSessionLocal = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False, ) async def get_db() -> AsyncIterator[AsyncSession]: """写路由用: 退出时自动 commit。""" async with AsyncSessionLocal() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() async def get_db_read() -> AsyncIterator[AsyncSession]: """读路由用: 不 commit(只读)。""" async with AsyncSessionLocal() as session: try: yield session finally: await session.close() async def init_db() -> None: """验证数据库连接(不建表)。 生产环境 schema 由 Alembic 管理。 本地开发/测试需要建表时调用 `create_all()`。 """ async with engine.begin() as conn: # 只验证连接,不自动建表 await conn.run_sync(lambda conn: None) async def create_all() -> None: """创建所有表(仅用于本地开发/测试)。""" async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all)