From f0903ac309d75b886e3e258fbfb0e11620773f99 Mon Sep 17 00:00:00 2001 From: CaptainB Date: Wed, 9 Jul 2025 11:02:55 +0800 Subject: [PATCH] refactor: enhance ToolTreeSerializer query parameters --- apps/tools/serializers/tool.py | 66 +++++++++++++++++----------------- apps/tools/views/tool.py | 45 ++++++++++++++--------- 2 files changed, 62 insertions(+), 49 deletions(-) diff --git a/apps/tools/serializers/tool.py b/apps/tools/serializers/tool.py index f2901a8cb..1323c28b4 100644 --- a/apps/tools/serializers/tool.py +++ b/apps/tools/serializers/tool.py @@ -16,7 +16,7 @@ from pylint.reporters import JSON2Reporter from rest_framework import serializers, status from common.database_model_manage.database_model_manage import DatabaseModelManage -from common.db.search import page_search, native_page_search +from common.db.search import page_search, native_page_search, native_search from common.exception.app_exception import AppApiException from common.field.common import UploadedImageField from common.result import result @@ -28,7 +28,6 @@ from maxkb.const import CONFIG, PROJECT_DIR from system_manage.models import AuthTargetType, WorkspaceUserResourcePermission from system_manage.serializers.user_resource_permission import UserResourcePermissionSerializer 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')) @@ -537,36 +536,6 @@ class ToolSerializer(serializers.Serializer): class ToolTreeSerializer(serializers.Serializer): - workspace_id = serializers.CharField(required=True, label=_('workspace id')) - scope = serializers.CharField(required=True, label=_('scope')) - - def get_tools(self, folder_id): - self.is_valid(raise_exception=True) - if not folder_id: - folder_id = self.data.get('workspace_id') - # 获取当前文件夹 - current_folder = ToolFolder.objects.filter(id=folder_id).first() - if not current_folder: - raise serializers.ValidationError(_('Folder not found')) - - # 获取当前文件夹下的直接子文件夹 - child_folders = ToolFolder.objects.filter(parent=current_folder) - folders_data = ToolFolderFlatSerializer(child_folders, many=True).data - - # 获取当前文件夹下的工具 - tools = QuerySet(Tool).filter( - Q(workspace_id=self.data.get('workspace_id')) & - Q(scope=self.data.get('scope')) & - Q(folder_id=folder_id) - ) - tools_data = ToolModelSerializer(tools, many=True).data - - # 返回包含文件夹和工具的结构 - return { - 'folders': folders_data, - 'tools': tools_data, - } - class Query(serializers.Serializer): workspace_id = serializers.CharField(required=True, label=_('workspace id')) folder_id = serializers.CharField(required=True, label=_('folder id')) @@ -677,3 +646,36 @@ class ToolTreeSerializer(serializers.Serializer): 'init_field_list': json.loads(record.get('init_field_list', '[]')), }, ) + + def get_tools(self): + 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() + results = native_search( + self.get_query_set(workspace_manage, is_x_pack_ee), + 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' + ) + ) + ), + + ) + + # 返回包含文件夹和工具的结构 + return { + 'folders': [ + folder for folder in results if folder['resource_type'] == 'folder' + ], + 'tools': [ + { + **tool, + 'input_field_list': json.loads(tool.get('input_field_list', '[]')), + 'init_field_list': json.loads(tool.get('init_field_list', '[]')), + } for tool in results if tool['resource_type'] == 'tool' + ], + } diff --git a/apps/tools/views/tool.py b/apps/tools/views/tool.py index ee0406675..f25cea7bd 100644 --- a/apps/tools/views/tool.py +++ b/apps/tools/views/tool.py @@ -67,9 +67,16 @@ class ToolView(APIView): RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), RoleConstants.USER.get_workspace_role() ) def get(self, request: Request, workspace_id: str): - return result.success(ToolTreeSerializer( - data={'workspace_id': workspace_id, 'scope': ToolScope.WORKSPACE, 'user_id': request.user.id} - ).get_tools(request.query_params.get('folder_id'))) + return result.success(ToolTreeSerializer.Query( + data={ + 'workspace_id': workspace_id, + 'folder_id': request.query_params.get('folder_id'), + 'name': request.query_params.get('name'), + 'scope': request.query_params.get('scope', ToolScope.WORKSPACE), + 'user_id': request.user.id, + 'create_user': request.query_params.get('create_user'), + } + ).get_tools()) class Debug(APIView): authentication_classes = [TokenAuth] @@ -111,8 +118,8 @@ class ToolView(APIView): PermissionConstants.TOOL_EDIT.get_workspace_permission_workspace_manage_role(), RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), ViewPermission([RoleConstants.USER.get_workspace_role()], - [PermissionConstants.TOOL.get_workspace_tool_permission()], - CompareConstants.AND), + [PermissionConstants.TOOL.get_workspace_tool_permission()], + CompareConstants.AND), ) @log( menu='Tool', operate='Update tool', @@ -136,9 +143,10 @@ class ToolView(APIView): @has_permissions( PermissionConstants.TOOL_READ.get_workspace_tool_permission(), PermissionConstants.TOOL_READ.get_workspace_permission_workspace_manage_role(), - RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), ViewPermission([RoleConstants.USER.get_workspace_role()], - [PermissionConstants.TOOL.get_workspace_tool_permission()], - CompareConstants.AND), + RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), + ViewPermission([RoleConstants.USER.get_workspace_role()], + [PermissionConstants.TOOL.get_workspace_tool_permission()], + CompareConstants.AND), ) @log(menu='Tool', operate='Get tool') def get(self, request: Request, workspace_id: str, tool_id: str): @@ -158,9 +166,10 @@ class ToolView(APIView): @has_permissions( PermissionConstants.TOOL_DELETE.get_workspace_tool_permission(), PermissionConstants.TOOL_DELETE.get_workspace_permission_workspace_manage_role(), - RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), ViewPermission([RoleConstants.USER.get_workspace_role()], - [PermissionConstants.TOOL.get_workspace_tool_permission()], - CompareConstants.AND), + RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), + ViewPermission([RoleConstants.USER.get_workspace_role()], + [PermissionConstants.TOOL.get_workspace_tool_permission()], + CompareConstants.AND), ) @log( menu='Tool', operate="Delete tool", @@ -242,9 +251,10 @@ class ToolView(APIView): @has_permissions( PermissionConstants.TOOL_EXPORT.get_workspace_tool_permission(), PermissionConstants.TOOL_EXPORT.get_workspace_permission_workspace_manage_role(), - RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), ViewPermission([RoleConstants.USER.get_workspace_role()], - [PermissionConstants.TOOL.get_workspace_tool_permission()], - CompareConstants.AND), + RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), + ViewPermission([RoleConstants.USER.get_workspace_role()], + [PermissionConstants.TOOL.get_workspace_tool_permission()], + CompareConstants.AND), ) @log( menu='Tool', operate="Export tool", @@ -296,9 +306,10 @@ class ToolView(APIView): @has_permissions( PermissionConstants.TOOL_EDIT.get_workspace_tool_permission(), PermissionConstants.TOOL_EDIT.get_workspace_permission_workspace_manage_role(), - RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), ViewPermission([RoleConstants.USER.get_workspace_role()], - [PermissionConstants.TOOL.get_workspace_tool_permission()], - CompareConstants.AND), + RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), + ViewPermission([RoleConstants.USER.get_workspace_role()], + [PermissionConstants.TOOL.get_workspace_tool_permission()], + CompareConstants.AND), ) def put(self, request: Request, tool_id: str, workspace_id: str): return result.success(ToolSerializer.IconOperate(data={