95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
from datetime import datetime
|
|
from sqlalchemy import event, true
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
|
|
from sqlalchemy.orm import Session, declarative_base, with_loader_criteria
|
|
from sqlalchemy.orm.attributes import flag_modified
|
|
from app.core.config import settings
|
|
from app.core.logging import get_logger
|
|
from typing import AsyncGenerator
|
|
|
|
from app.db.datascope import DataScopeHandle, SYSTEM_USER
|
|
from app.db.models.base_entity import BaseEntity
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
# 创建数据库引擎
|
|
engine = create_async_engine(
|
|
settings.database_url,
|
|
echo=False, # 关闭SQL调试日志以减少输出
|
|
future=True,
|
|
pool_pre_ping=True,
|
|
pool_recycle=3600
|
|
)
|
|
|
|
# 创建会话工厂
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False
|
|
)
|
|
|
|
@event.listens_for(Session, "do_orm_execute")
|
|
def _apply_data_scope(orm_execute_state):
|
|
# only act on SELECT queries generated by ORM
|
|
if not orm_execute_state.is_select:
|
|
return
|
|
|
|
allowed = DataScopeHandle.allowed_users()
|
|
if not allowed or len(allowed) == 0:
|
|
return
|
|
|
|
# predicate builder: return None to skip for classes that opt-out
|
|
def criteria_fn(cls):
|
|
# skip if the mapped class explicitly disables data-scope
|
|
if getattr(cls, "__ignore_data_scope__", False):
|
|
return true()
|
|
# some classes may not have created_by column; guard dynamically
|
|
col = getattr(cls, "created_by", None)
|
|
if col is None:
|
|
return true()
|
|
return col.in_(allowed)
|
|
|
|
# apply loader-level criteria to all subclasses of BaseEntity
|
|
orm_execute_state.statement = orm_execute_state.statement.options(
|
|
with_loader_criteria(BaseEntity, criteria_fn, include_aliases=True)
|
|
)
|
|
|
|
@event.listens_for(Session, "before_flush")
|
|
def _audit_before_flush(session, flush_context, instances):
|
|
user = DataScopeHandle.get_user_info()
|
|
effective_user = (user or "").strip() or SYSTEM_USER
|
|
now = datetime.now()
|
|
|
|
# new -> set created_* and updated_*
|
|
for obj in list(session.new):
|
|
if isinstance(obj, BaseEntity):
|
|
if getattr(obj, "created_at", None) is None:
|
|
obj.created_at = now
|
|
if getattr(obj, "created_by", None) is None:
|
|
obj.created_by = effective_user
|
|
obj.updated_at = now
|
|
obj.updated_by = effective_user
|
|
# ensure SQLAlchemy sees changes
|
|
try:
|
|
flag_modified(obj, "created_by")
|
|
except Exception:
|
|
pass
|
|
|
|
# dirty -> set updated_*
|
|
for obj in list(session.dirty):
|
|
if isinstance(obj, BaseEntity):
|
|
obj.updated_at = now
|
|
obj.updated_by = effective_user
|
|
try:
|
|
flag_modified(obj, "updated_by")
|
|
except Exception:
|
|
pass
|
|
|
|
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
|
"""获取数据库会话"""
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
yield session
|
|
finally:
|
|
await session.close()
|