feat: enhance tool query logic and add user-specific SQL files for improved access control

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_query_set = QuerySet(Tool).filter(workspace_id=self.data.get('workspace_id'))
-            tool_scope_query_set = QuerySet(Tool)
             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)
+                default_query_set = default_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)
             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)
+                default_query_set = default_query_set.filter(desc__contains=desc)
-            tool_query_set = tool_query_set.order_by("-update_time")
+
+            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', '[]')),

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}
+     ${default_query_set}

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}

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}

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))