From 8fc074fecbf277ecd9ac2b43d95bcb135e83d009 Mon Sep 17 00:00:00 2001 From: CaptainB Date: Thu, 26 Jun 2025 19:05:33 +0800 Subject: [PATCH] feat: enhance tool query logic and add user-specific SQL files for improved access control --- apps/tools/serializers/tool.py | 51 +++++++++++++++++------- apps/tools/sql/list_tool.sql | 4 +- apps/tools/sql/list_tool_user.sql | 49 +++++++++++++++++++++++ apps/tools/sql/list_tool_user_ee.sql | 59 ++++++++++++++++++++++++++++ apps/tools/views/tool.py | 1 + 5 files changed, 148 insertions(+), 16 deletions(-) create mode 100644 apps/tools/sql/list_tool_user.sql create mode 100644 apps/tools/sql/list_tool_user_ee.sql diff --git a/apps/tools/serializers/tool.py b/apps/tools/serializers/tool.py index c8e0d93f5..323d95f9c 100644 --- a/apps/tools/serializers/tool.py +++ b/apps/tools/serializers/tool.py @@ -18,6 +18,7 @@ from rest_framework import serializers, status from common.constants.cache_version import Cache_Version from common.constants.permission_constants import ResourceAuthType, ResourcePermissionGroup +from common.database_model_manage.database_model_manage import DatabaseModelManage from common.db.search import page_search, native_page_search from common.exception.app_exception import AppApiException from common.field.common import UploadedImageField @@ -30,6 +31,7 @@ from maxkb.const import CONFIG, PROJECT_DIR from system_manage.models import AuthTargetType, WorkspaceUserResourcePermission from tools.models import Tool, ToolScope, ToolFolder, ToolType from tools.serializers.tool_folder import ToolFolderFlatSerializer +from users.serializers.user import is_workspace_manage tool_executor = ToolExecutor(CONFIG.get('SANDBOX')) @@ -541,7 +543,7 @@ class ToolTreeSerializer(serializers.Serializer): workspace_id = serializers.CharField(required=True, label=_('workspace id')) folder_id = serializers.CharField(required=True, label=_('folder id')) name = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool name')) - user_id = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('user id')) + user_id = serializers.UUIDField(required=False, allow_null=True, label=_('user id')) scope = serializers.CharField(required=True, label=_('scope')) def page_tool(self, current_page: int, page_size: int): @@ -570,9 +572,10 @@ class ToolTreeSerializer(serializers.Serializer): return page_search(current_page, page_size, tools, lambda record: ToolModelSerializer(record).data) def get_query_set(self): - tool_query_set = QuerySet(Tool) - tool_scope_query_set = QuerySet(Tool) + tool_query_set = QuerySet(Tool).filter(workspace_id=self.data.get('workspace_id')) folder_query_set = QuerySet(ToolFolder) + default_query_set = QuerySet(Tool) + workspace_id = self.data.get('workspace_id') user_id = self.data.get('user_id') scope = self.data.get('scope') @@ -582,36 +585,56 @@ class ToolTreeSerializer(serializers.Serializer): if workspace_id is not None: folder_query_set = folder_query_set.filter(workspace_id=workspace_id) - tool_query_set = tool_query_set.filter(workspace_id=workspace_id) - if user_id is not None: - folder_query_set = folder_query_set.filter(user_id=user_id) - tool_query_set = tool_query_set.filter(user_id=user_id) + default_query_set = default_query_set.filter(workspace_id=workspace_id) if folder_id is not None: folder_query_set = folder_query_set.filter(parent=folder_id) - tool_query_set = tool_query_set.filter(folder_id=folder_id) + default_query_set = default_query_set.filter(folder_id=folder_id) if name is not None: folder_query_set = folder_query_set.filter(name__contains=name) - tool_query_set = tool_query_set.filter(name__contains=name) + default_query_set = default_query_set.filter(name__contains=name) if desc is not None: folder_query_set = folder_query_set.filter(desc__contains=desc) - tool_query_set = tool_query_set.filter(desc__contains=desc) - tool_query_set = tool_query_set.order_by("-update_time") + default_query_set = default_query_set.filter(desc__contains=desc) + + default_query_set = default_query_set.order_by("-create_time") if scope is not None: - tool_scope_query_set = tool_scope_query_set.filter(scope=scope) + tool_query_set = tool_query_set.filter(scope=scope) return { 'folder_query_set': folder_query_set, 'tool_query_set': tool_query_set, - 'tool_scope_query_set': tool_scope_query_set + 'default_query_set': default_query_set, + 'workspace_user_resource_permission_query_set': QuerySet(WorkspaceUserResourcePermission).filter( + auth_target_type="TOOL", + workspace_id=workspace_id, + user_id=user_id + ) } + @staticmethod + def is_x_pack_ee(): + workspace_user_role_mapping_model = DatabaseModelManage.get_model("workspace_user_role_mapping") + role_permission_mapping_model = DatabaseModelManage.get_model("role_permission_mapping_model") + return workspace_user_role_mapping_model is not None and role_permission_mapping_model is not None + def page_tool_with_folders(self, current_page: int, page_size: int): self.is_valid(raise_exception=True) + workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id')) + is_x_pack_ee = self.is_x_pack_ee() + return native_page_search( current_page, page_size, self.get_query_set(), - get_file_content(os.path.join(PROJECT_DIR, "apps", "tools", 'sql', 'list_tool.sql')), + get_file_content( + os.path.join( + PROJECT_DIR, + "apps", "tools", 'sql', + 'list_tool.sql' if workspace_manage else ( + 'list_tool_user_ee.sql' if is_x_pack_ee else 'list_tool_user.sql' + ) + ) + ), post_records_handler=lambda record: { **record, 'input_field_list': json.loads(record.get('input_field_list', '[]')), diff --git a/apps/tools/sql/list_tool.sql b/apps/tools/sql/list_tool.sql index 21715ed9a..978042813 100644 --- a/apps/tools/sql/list_tool.sql +++ b/apps/tools/sql/list_tool.sql @@ -18,7 +18,7 @@ from (select tool."id"::text, tool.input_field_list, tool."is_active" from tool - left join "user" on "user".id = user_id ${tool_scope_query_set} + left join "user" on "user".id = user_id ${tool_query_set} UNION select tool_folder."id", tool_folder."name", @@ -40,4 +40,4 @@ from (select tool."id"::text, 'true' as "is_active" from tool_folder left join "user" on "user".id = user_id ${folder_query_set}) temp - ${tool_query_set} \ No newline at end of file + ${default_query_set} \ No newline at end of file diff --git a/apps/tools/sql/list_tool_user.sql b/apps/tools/sql/list_tool_user.sql new file mode 100644 index 000000000..07c1c471a --- /dev/null +++ b/apps/tools/sql/list_tool_user.sql @@ -0,0 +1,49 @@ +SELECT * +FROM (SELECT tool."id"::text, + tool."name", + tool."desc", + tool."tool_type", + tool."scope", + 'tool' AS "resource_type", + tool."workspace_id", + tool."folder_id", + tool."user_id", + "user".nick_name AS "nick_name", + tool."icon", + tool.label, + tool."template_id"::text, + tool."create_time", + tool."update_time", + tool.init_field_list, + tool.input_field_list, + tool."is_active" + FROM (SELECT tool.* + FROM tool tool ${tool_query_set} + AND tool.id IN (SELECT target + FROM workspace_user_resource_permission + WHERE auth_target_type = 'TOOL' + AND 'VIEW' = ANY (permission_list))) AS tool + LEFT JOIN "user" ON "user".id = user_id + + UNION + SELECT tool_folder."id", + tool_folder."name", + tool_folder."desc", + 'folder' AS "tool_type", + '' AS scope, + 'folder' AS "resource_type", + tool_folder."workspace_id", + tool_folder."parent_id" AS "folder_id", + tool_folder."user_id", + "user".nick_name AS "nick_name", + '' AS "icon", + '' AS label, + '' AS "template_id", + tool_folder."create_time", + tool_folder."update_time", + '[]'::jsonb AS init_field_list, + '[]'::jsonb AS input_field_list, + 'true' AS "is_active" + FROM tool_folder + LEFT JOIN "user" ON "user".id = user_id ${folder_query_set}) temp + ${default_query_set} \ No newline at end of file diff --git a/apps/tools/sql/list_tool_user_ee.sql b/apps/tools/sql/list_tool_user_ee.sql new file mode 100644 index 000000000..00c7b808a --- /dev/null +++ b/apps/tools/sql/list_tool_user_ee.sql @@ -0,0 +1,59 @@ +SELECT * +FROM (SELECT tool."id"::text, + tool."name", + tool."desc", + tool."tool_type", + tool."scope", + 'tool' AS "resource_type", + tool."workspace_id", + tool."folder_id", + tool."user_id", + "user".nick_name AS "nick_name", + tool."icon", + tool.label, + tool."template_id"::text, + tool."create_time", + tool."update_time", + tool.init_field_list, + tool.input_field_list, + tool."is_active" + FROM (SELECT tool.* + FROM tool tool ${tool_query_set} + AND tool.id IN (SELECT target + FROM workspace_user_resource_permission ${workspace_user_resource_permission_query_set} + AND CASE + WHEN auth_type = 'ROLE' THEN + 'ROLE' = ANY (permission_list) + AND + 'TOOL:READ' IN (SELECT (CASE WHEN user_role_relation.role_id = ANY (ARRAY ['USER']) THEN 'TOOL:READ' ELSE role_permission.permission_id END) + FROM role_permission role_permission + RIGHT JOIN user_role_relation user_role_relation ON user_role_relation.role_id=role_permission.role_id + WHERE user_role_relation.user_id=workspace_user_resource_permission.user_id + AND user_role_relation.workspace_id=workspace_user_resource_permission.workspace_id) + ELSE + 'VIEW' = ANY (permission_list) + END + )) AS tool + LEFT JOIN "user" ON "user".id = user_id + UNION + SELECT tool_folder."id", + tool_folder."name", + tool_folder."desc", + 'folder' AS "tool_type", + '' AS scope, + 'folder' AS "resource_type", + tool_folder."workspace_id", + tool_folder."parent_id" AS "folder_id", + tool_folder."user_id", + "user".nick_name AS "nick_name", + '' AS "icon", + '' AS label, + '' AS "template_id", + tool_folder."create_time", + tool_folder."update_time", + '[]'::jsonb AS init_field_list, + '[]'::jsonb AS input_field_list, + 'true' AS "is_active" + FROM tool_folder + LEFT JOIN "user" ON "user".id = user_id ${folder_query_set}) temp + ${default_query_set} \ No newline at end of file diff --git a/apps/tools/views/tool.py b/apps/tools/views/tool.py index fed20c0a8..f2680a0dc 100644 --- a/apps/tools/views/tool.py +++ b/apps/tools/views/tool.py @@ -183,6 +183,7 @@ class ToolView(APIView): 'folder_id': request.query_params.get('folder_id'), 'name': request.query_params.get('name'), 'scope': request.query_params.get('scope'), + 'user_id': request.user.id } ).page_tool_with_folders(current_page, page_size))