dibang-xiaoshou/backend/app/db/session.py

from collections.abc import AsyncGenerator

from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession

from app.core.config import get_settings

# Import all models to ensure they are registered with SQLModel.metadata
from app.models import (  # noqa: F401
    Cargo,
    Company,
    Inventory,
    Plan,
    User,
    Vehicle,
    Weighing,
)

settings = get_settings()

# 配置连接池参数
# SQLAlchemy 内置连接池，无需额外插件
engine = create_async_engine(
    settings.database_url,
    echo=settings.env == "dev",
    # 连接池配置（从环境变量读取）
    pool_size=settings.db_pool_size,
    max_overflow=settings.db_max_overflow,
    pool_timeout=settings.db_pool_timeout,
    pool_recycle=settings.db_pool_recycle,
    pool_pre_ping=settings.db_pool_pre_ping,
)

AsyncSessionLocal = async_sessionmaker(
    bind=engine, class_=AsyncSession, expire_on_commit=False, autoflush=False
)


async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSessionLocal() as session:
        yield session


async def init_models() -> None:
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)