deepdata: core product code (backend, frontend, runtime, deepflow module)

runtime/deepdata-python/app/db/session.py

@@ -0,0 +1,94 @@
+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()