101 lines
2.4 KiB
Python
101 lines
2.4 KiB
Python
"""
|
||
数据库会话管理
|
||
"""
|
||
from typing import AsyncGenerator
|
||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||
from sqlalchemy import create_engine
|
||
from sqlalchemy.orm import sessionmaker, Session
|
||
from app.core.config import settings
|
||
from app.db.base import Base
|
||
|
||
# 创建异步引擎
|
||
engine = create_async_engine(
|
||
settings.DATABASE_URL,
|
||
echo=settings.DATABASE_ECHO,
|
||
pool_pre_ping=True,
|
||
pool_size=50, # 从20增加到50，提高并发性能
|
||
max_overflow=10, # 从0增加到10，允许峰值时的额外连接
|
||
)
|
||
|
||
# 创建同步引擎（用于遗留同步查询）
|
||
def _get_sync_database_url() -> str:
|
||
url = settings.DATABASE_URL
|
||
if url.startswith("postgresql+asyncpg://"):
|
||
return url.replace("postgresql+asyncpg://", "postgresql+psycopg2://", 1)
|
||
if "+asyncpg" in url:
|
||
return url.replace("+asyncpg", "+psycopg2")
|
||
return url
|
||
|
||
|
||
sync_engine = create_engine(
|
||
_get_sync_database_url(),
|
||
echo=settings.DATABASE_ECHO,
|
||
pool_pre_ping=True,
|
||
pool_size=50,
|
||
max_overflow=10,
|
||
)
|
||
|
||
# 创建异步会话工厂
|
||
async_session_maker = async_sessionmaker(
|
||
engine,
|
||
class_=AsyncSession,
|
||
expire_on_commit=False,
|
||
autocommit=False,
|
||
autoflush=False,
|
||
)
|
||
|
||
# 创建同步会话工厂
|
||
sync_session_maker = sessionmaker(
|
||
bind=sync_engine,
|
||
autocommit=False,
|
||
autoflush=False,
|
||
)
|
||
|
||
|
||
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||
"""
|
||
获取数据库会话
|
||
|
||
Yields:
|
||
AsyncSession: 数据库会话
|
||
"""
|
||
async with async_session_maker() as session:
|
||
try:
|
||
yield session
|
||
await session.commit()
|
||
except Exception:
|
||
await session.rollback()
|
||
raise
|
||
finally:
|
||
await session.close()
|
||
|
||
|
||
async def init_db() -> None:
|
||
"""
|
||
初始化数据库（创建所有表）
|
||
注意：生产环境应使用Alembic迁移
|
||
"""
|
||
async with engine.begin() as conn:
|
||
# 导入所有模型以确保它们被注册
|
||
from app.models import user, asset, device_type, organization
|
||
|
||
# 创建所有表
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
|
||
|
||
async def close_db() -> None:
|
||
"""关闭数据库连接"""
|
||
await engine.dispose()
|
||
sync_engine.dispose()
|
||
|
||
|
||
__all__ = [
|
||
"engine",
|
||
"sync_engine",
|
||
"async_session_maker",
|
||
"sync_session_maker",
|
||
"get_db",
|
||
"init_db",
|
||
"close_db",
|
||
]
|