74 lines
1.5 KiB
Python
74 lines
1.5 KiB
Python
"""
|
|
数据库配置和会话管理
|
|
使用SQLAlchemy 2.x异步支持
|
|
"""
|
|
from typing import AsyncGenerator
|
|
|
|
from sqlalchemy import MetaData
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from app.config import settings
|
|
|
|
# 创建数据库引擎
|
|
# 使用NullPool避免连接池在多进程环境中的问题
|
|
engine = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
echo=settings.DB_ECHO,
|
|
poolclass=NullPool,
|
|
future=True,
|
|
)
|
|
|
|
# 创建异步会话工厂
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
bind=engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autoflush=True,
|
|
autocommit=False,
|
|
)
|
|
|
|
# 创建基类
|
|
Base = declarative_base()
|
|
|
|
# 元数据（用于迁移）
|
|
metadata = MetaData()
|
|
|
|
|
|
async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
|
|
"""
|
|
获取异步数据库会话
|
|
使用依赖注入模式
|
|
"""
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
async def init_db() -> None:
|
|
"""
|
|
初始化数据库
|
|
创建所有表
|
|
"""
|
|
async with engine.begin() as conn:
|
|
# 创建所有表
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
|
|
async def close_db() -> None:
|
|
"""
|
|
关闭数据库连接
|
|
"""
|
|
await engine.dispose()
|