| 1234567891011121314151617181920212223242526272829303132333435363738394041424344 |
- from sqlalchemy import create_engine, event
- from sqlalchemy.ext.declarative import declarative_base
- from sqlalchemy.orm import sessionmaker
- from app.core.config import settings
- engine = create_engine(
- settings.DATABASE_URL,
- pool_pre_ping=True,
- pool_recycle=settings.DB_POOL_RECYCLE_SEC,
- pool_size=settings.DB_POOL_SIZE,
- max_overflow=settings.DB_MAX_OVERFLOW,
- pool_reset_on_return="rollback",
- connect_args={
- "connect_timeout": settings.DB_CONNECT_TIMEOUT_SEC,
- "read_timeout": settings.DB_READ_TIMEOUT_SEC,
- "write_timeout": settings.DB_WRITE_TIMEOUT_SEC,
- },
- echo=settings.DEBUG,
- )
- # 避免 commit 后 ORM 实例自动过期，导致读取字段时回表 SELECT（高频接口会显著拖慢）
- SessionLocal = sessionmaker(autocommit=False, autoflush=False, expire_on_commit=False, bind=engine)
- Base = declarative_base()
- def get_db():
- """获取数据库会话"""
- db = SessionLocal()
- try:
- yield db
- finally:
- db.close()
- @event.listens_for(Base, "before_insert", propagate=True)
- def _auto_fill_create_dept_id(mapper, connection, target):
- """ORM before_insert 钩子：当 create_dept_id 为空时，自动从请求上下文填充当前用户部门 ID"""
- if hasattr(target, "create_dept_id") and not getattr(target, "create_dept_id", None):
- from app.core.context import current_user_dept_id
- dept_id = current_user_dept_id.get("")
- if dept_id:
- target.create_dept_id = dept_id
|
