From 5fd5ca258beec0472c08886c5309295df24dad9c Mon Sep 17 00:00:00 2001 From: hhhhsc <1710496817@qq.com> Date: Thu, 27 Aug 2026 09:30:59 +0800 Subject: [PATCH] =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E6=9F=A5=E8=AF=A2=E5=85=BC?= =?UTF-8?q?=E5=AE=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- runtime/datamate-python/app/db/datascope.py | 5 ++ .../datamate-python/app/db/models/models.py | 3 +- runtime/datamate-python/app/db/session.py | 30 ++++++++---- runtime/datamate-python/app/middleware.py | 1 - .../module/system/service/models_service.py | 32 ++++++++++--- .../tests/test_module_system.py | 46 +++++++++++++++++++ 6 files changed, 98 insertions(+), 19 deletions(-) diff --git a/runtime/datamate-python/app/db/datascope.py b/runtime/datamate-python/app/db/datascope.py index a8d19541f..eac05ced1 100644 --- a/runtime/datamate-python/app/db/datascope.py +++ b/runtime/datamate-python/app/db/datascope.py @@ -3,6 +3,7 @@ _current_user: ContextVar[Optional[str]] = ContextVar("_current_user", default=None) SYSTEM_USER = "system" +ADMIN_USER = "admin" class DataScopeHandle: @@ -26,6 +27,10 @@ def remove_user_info() -> None: def get_user_info() -> Optional[str]: return _current_user.get() + @staticmethod + def is_admin() -> bool: + return DataScopeHandle.get_user_info() == ADMIN_USER + @staticmethod def allowed_users() -> List[str]: """ diff --git a/runtime/datamate-python/app/db/models/models.py b/runtime/datamate-python/app/db/models/models.py index efec5506f..5c0a691fd 100644 --- a/runtime/datamate-python/app/db/models/models.py +++ b/runtime/datamate-python/app/db/models/models.py @@ -4,7 +4,7 @@ class Models(BaseEntity): - """模型配置表，对应表 t_models。模型为系统级配置，RAG/生成等按 ID 引用时不受数据权限过滤。 + """模型配置表，对应表 t_models。普通用户按创建者隔离，管理员可跨用户管理。 CREATE TABLE IF NOT EXISTS t_models ( id VARCHAR(36) PRIMARY KEY COMMENT '主键ID', @@ -24,6 +24,7 @@ class Models(BaseEntity): """ __tablename__ = "t_models" + __admin_ignore_data_scope__ = True id = Column(String(36), primary_key=True, index=True, comment="主键ID") model_name = Column(String(100), nullable=False, comment="模型名称（如 qwen2）") diff --git a/runtime/datamate-python/app/db/session.py b/runtime/datamate-python/app/db/session.py index a4d9a9304..3b52a772e 100644 --- a/runtime/datamate-python/app/db/session.py +++ b/runtime/datamate-python/app/db/session.py @@ -37,17 +37,27 @@ def _apply_data_scope(orm_execute_state): allowed = DataScopeHandle.allowed_users() if not allowed or len(allowed) == 0: return + is_admin = DataScopeHandle.is_admin() - # 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) + if is_admin: + def criteria_fn(cls): + if ( + getattr(cls, "__ignore_data_scope__", False) + or getattr(cls, "__admin_ignore_data_scope__", False) + ): + return true() + col = getattr(cls, "created_by", None) + if col is None: + return true() + return col.in_(allowed) + else: + def criteria_fn(cls): + if getattr(cls, "__ignore_data_scope__", False): + return true() + 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( diff --git a/runtime/datamate-python/app/middleware.py b/runtime/datamate-python/app/middleware.py index 5a6aac53f..64fd616d7 100644 --- a/runtime/datamate-python/app/middleware.py +++ b/runtime/datamate-python/app/middleware.py @@ -50,4 +50,3 @@ async def dispatch(self, request: Request, call_next): # Re-export ExceptionHandlingMiddleware for backward compatibility __all__ = ['UserContextMiddleware'] - diff --git a/runtime/datamate-python/app/module/system/service/models_service.py b/runtime/datamate-python/app/module/system/service/models_service.py index e093e836b..55c190149 100644 --- a/runtime/datamate-python/app/module/system/service/models_service.py +++ b/runtime/datamate-python/app/module/system/service/models_service.py @@ -35,6 +35,13 @@ ] +def _created_by_match(user: str): + return ( + (Models.created_by == user) + | ((Models.created_by.is_(None)) & (user == SYSTEM_USER)) + ) + + def _orm_to_response(row: Models) -> ModelsResponse: return ModelsResponse( id=row.id, @@ -115,10 +122,7 @@ async def create_model(self, req: CreateModelRequest) -> ModelsResponse: # 检查 uk_model_provider (model_name, provider, created_by) 是否存在且已删除 effective_user = (DataScopeHandle.get_user_info() or "").strip() or SYSTEM_USER - created_by_match = ( - (Models.created_by == effective_user) - | ((Models.created_by.is_(None)) & (effective_user == SYSTEM_USER)) - ) + created_by_match = _created_by_match(effective_user) deleted_row = ( await self.db.execute( select(Models).where( @@ -143,7 +147,11 @@ async def create_model(self, req: CreateModelRequest) -> ModelsResponse: if want_default: await self.db.execute( update(Models) - .where(Models.type == req.type.value, Models.is_default == True) + .where( + Models.type == req.type.value, + Models.is_default == True, + created_by_match, + ) .values(is_default=False) ) entity.is_default = want_default @@ -160,6 +168,7 @@ async def create_model(self, req: CreateModelRequest) -> ModelsResponse: (Models.is_deleted == False) | (Models.is_deleted.is_(None)), Models.type == req.type.value, Models.is_default == True, + created_by_match, ) ) ).scalar_one_or_none() @@ -170,7 +179,11 @@ async def create_model(self, req: CreateModelRequest) -> ModelsResponse: else: await self.db.execute( update(Models) - .where(Models.type == req.type.value, Models.is_default == True) + .where( + Models.type == req.type.value, + Models.is_default == True, + created_by_match, + ) .values(is_default=False) ) is_default = req.isDefault if req.isDefault is not None else False @@ -220,9 +233,14 @@ async def update_model(self, model_id: str, req: CreateModelRequest) -> ModelsRe want_default = req.isDefault if req.isDefault is not None else False if (entity.is_default is not True) and want_default: + owner = (entity.created_by or "").strip() or SYSTEM_USER await self.db.execute( update(Models) - .where(Models.type == req.type.value, Models.is_default == True) + .where( + Models.type == req.type.value, + Models.is_default == True, + _created_by_match(owner), + ) .values(is_default=False) ) entity.is_default = want_default diff --git a/runtime/datamate-python/tests/test_module_system.py b/runtime/datamate-python/tests/test_module_system.py index 01349fef1..9c65464a6 100644 --- a/runtime/datamate-python/tests/test_module_system.py +++ b/runtime/datamate-python/tests/test_module_system.py @@ -2,6 +2,12 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock +from sqlalchemy import select + +from app.db.datascope import DataScopeHandle +from app.db.models.dataset_management import Dataset +from app.db.models.models import Models +from app.db.session import _apply_data_scope from app.module.system.service.common_service import get_model_by_id @@ -65,3 +71,43 @@ def test_get_model_by_id_returns_exact_scalar_object() -> None: fetched = _run(get_model_by_id(db, "m100")) assert fetched is model_obj + + +def _compile_scoped_query(model, user: str) -> str: + DataScopeHandle.set_user_info(user) + try: + state = SimpleNamespace(is_select=True, statement=select(model)) + _apply_data_scope(state) + return str(state.statement.compile(compile_kwargs={"literal_binds": True})) + finally: + DataScopeHandle.remove_user_info() + + +def test_regular_user_model_query_keeps_creator_scope() -> None: + sql = _compile_scoped_query(Models, "alice") + + assert "t_models.created_by IN ('alice', 'system')" in sql + + +def test_admin_model_query_bypasses_creator_scope() -> None: + sql = _compile_scoped_query(Models, "admin") + + assert "t_models.created_by IN" not in sql + + +def test_similar_username_does_not_gain_admin_model_access() -> None: + sql = _compile_scoped_query(Models, "Admin") + + assert "t_models.created_by IN ('Admin', 'system')" in sql + + +def test_system_user_does_not_gain_admin_model_access() -> None: + sql = _compile_scoped_query(Models, "system") + + assert "t_models.created_by IN ('system', 'system')" in sql + + +def test_admin_query_keeps_scope_for_unmarked_models() -> None: + sql = _compile_scoped_query(Dataset, "admin") + + assert "t_dm_datasets.created_by IN ('admin', 'system')" in sql