Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 25 additions & 9 deletions apps/application/serializers/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
from tools.models import Tool, ToolScope
from tools.serializers.tool import ToolModelSerializer
from users.models import User
from users.serializers.user import is_workspace_manage


def get_base_node_work_flow(work_flow):
Expand Down Expand Up @@ -290,9 +291,10 @@ class ApplicationListResponse(serializers.Serializer):


class Query(serializers.Serializer):
workspace_id = serializers.CharField(required=False, label=_('workspace id'))
workspace_id = serializers.CharField(required=False, label=_('Workspace ID'))
user_id = serializers.UUIDField(required=True, label=_("User ID"))

def get_query_set(self, instance: Dict):
def get_query_set(self, instance: Dict, workspace_manage: bool):
folder_query_set = QuerySet(ApplicationFolder)
application_query_set = QuerySet(Application)
workspace_id = self.data.get('workspace_id')
Expand All @@ -315,11 +317,14 @@ def get_query_set(self, instance: Dict):
if desc is not None:
folder_query_set = folder_query_set.filter(desc__contains=desc)
application_query_set = application_query_set.filter(desc__contains=desc)
application_custom_sql_query_set = application_query_set
application_query_set = application_query_set.order_by("-update_time")
return {
'folder_query_set': folder_query_set,
'application_query_set': application_query_set
}
'application_query_set': application_query_set,
'application_custom_sql': application_custom_sql_query_set
} if workspace_manage else {'folder_query_set': folder_query_set,
'application_query_set': application_query_set}

@staticmethod
def is_x_pack_ee():
Expand All @@ -329,17 +334,28 @@ def is_x_pack_ee():

def list(self, instance: Dict):
self.is_valid(raise_exception=True)
workspace_id = self.data.get('workspace_id')
user_id = self.data.get("user_id")
ApplicationQueryRequest(data=instance).is_valid(raise_exception=True)
return native_search(self.get_query_set(instance), select_string=get_file_content(
workspace_manage = is_workspace_manage(user_id, workspace_id)

return native_search(self.get_query_set(instance, workspace_manage), select_string=get_file_content(
os.path.join(PROJECT_DIR, "apps", "application", 'sql',
'list_application_ee.sql' if self.is_x_pack_ee() else 'list_application.sql')))
'list_application.sql' if workspace_manage else (
'list_application_user_ee.sql' if self.is_x_pack_ee() else 'list_application_user.sql')
)))

def page(self, current_page: int, page_size: int, instance: Dict):
self.is_valid(raise_exception=True)
ApplicationQueryRequest(data=instance).is_valid(raise_exception=True)
return native_page_search(current_page, page_size, self.get_query_set(instance), get_file_content(
os.path.join(PROJECT_DIR, "apps", "application", 'sql',
'list_application_ee.sql' if self.is_x_pack_ee() else 'list_application.sql')),
workspace_id = self.data.get('workspace_id')
user_id = self.data.get("user_id")
workspace_manage = is_workspace_manage(user_id, workspace_id)
return native_page_search(current_page, page_size, self.get_query_set(instance, workspace_manage),
get_file_content(
os.path.join(PROJECT_DIR, "apps", "application", 'sql',
'list_application.sql' if workspace_manage else (
'list_application_user_ee.sql' if self.is_x_pack_ee() else 'list_application_user.sql'))),
)


Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here are some suggestions to improve your code:

  1. Variable Naming: Replace instance: Dict with something more specific like request_data to describe its purpose.

  2. Comments and Docstrings:

    • Add comments on complex operations to clarify what they do.
    • Use docstrings for functions and methods where applicable to ensure clarity.
  3. Separation of concerns:

    • Create utility functions that encapsulate logic rather than mixing it within classes or methods.
  4. Error Handling:

    • Ensure exceptions are handled appropriately to avoid silent failures.
  5. Security Considerations:

    • Validate and sanitize input before using them in queries.
  6. Code Organization:

    • Group related functionality together within the same class or file for better readability.

Here's an improved version based on these points:

# Import necessary modules
from rest_framework import serializers
import os

# Path configuration
PROJECT_DIR = '/path/to/project'

def is_workspace_manage(user_uuid, workspace_id):
    # Implement the function to check if user has manage permissions for the workflow
    pass

class BaseNodeWorkFlowViewset(ViewSetMixin, GenericViewSet):

    ...

    def get_query_set(self, request_data: dict, include_custom_queries=True) -> tuple[QuerySet[Any], QuerySet[Any]]:
        """Return query sets filtered by workspace."""
        folder_qs = QuerySet(ApplicationFolder).filter(workspace=request_data['workspace'])
        app_qs = QuerySet(Application).filter(workspace=request_data['workspace'])

        desc = request_data.get('description', '')
        for qs in (folder_qs, app_qs):
            if desc:
                qs = qs.filter(description__icontains=desc)

        sorted_app_qs = app_qs.order_by('-updated_at')

        custom_app_qs = app_qs

        if not include_custom_queries:
            custom_app_qs = []

        return folder_qs, sorted_app_qs, custom_app_qs

    async def retrieve(self, request):
        try:
            data = await self.serializer_class.validate(request.data)
            return await native_page_search(page_num=data.page,
                                         page_size=data.size,
                                         query=self.get_query_set(data))
        except ValidationError as e:
            http_error_handler(e.detail)

Key Changes:

  1. Functionality Encapsulation: Created a method get_query_set() which simplifies handling query logic outside serialization methods.

  2. Input Validation: The retrieve() method includes validation to handle incoming requests correctly.

  3. Error Handling: Added error handling using asyncio.ErrorHandler.

  4. Documentation: Added docstrings for methods to explain their parameters and responsibilities.

This refactored approach makes the code cleaner, easier to understand, and maintainable.

Expand Down
5 changes: 1 addition & 4 deletions apps/application/sql/list_application.sql
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,7 @@ from (select "id"::text,
"create_time",
"update_time"
from application
where id in (select target
from workspace_user_resource_permission
where auth_target_type = 'APPLICATION'
and 'VIEW' = any (permission_list))
${application_custom_sql}
UNION
select "id",
"name",
Expand Down
31 changes: 31 additions & 0 deletions apps/application/sql/list_application_user.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
select *
from (select "id"::text,
"name",
"desc",
"is_publish",
"type",
'application' as "resource_type",
"workspace_id",
"folder_id",
"user_id",
"create_time",
"update_time"
from application
where id in (select target
from workspace_user_resource_permission
where auth_target_type = 'APPLICATION'
and 'VIEW' = any (permission_list))
UNION
select "id",
"name",
"desc",
true as "is_publish",
'folder' as "type",
'folder' as "resource_type",
"workspace_id",
"parent_id" as "folder_id",
"user_id",
"create_time",
"update_time"
from application_folder ${folder_query_set}) temp
${application_query_set}
5 changes: 3 additions & 2 deletions apps/application/views/application_chat_record.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from common import result
from common.auth import TokenAuth
from common.auth.authentication import has_permissions
from common.constants.permission_constants import PermissionConstants
from common.constants.permission_constants import PermissionConstants, RoleConstants
from common.utils.common import query_params_to_single_dict


Expand All @@ -35,7 +35,8 @@ class ApplicationChatRecord(APIView):
responses=ApplicationChatRecordQueryAPI.get_response(),
tags=[_("Application/Conversation Log")] # type: ignore
)
@has_permissions(PermissionConstants.APPLICATION_CHAT_LOG.get_workspace_application_permission())
@has_permissions(PermissionConstants.APPLICATION_CHAT_LOG.get_workspace_application_permission(),
RoleConstants.WORKSPACE_MANAGE.get_workspace_role())
def get(self, request: Request, workspace_id: str, application_id: str, chat_id: str):
return result.success(ApplicationChatRecordQuerySerializers(
data={**query_params_to_single_dict(request.query_params), 'application_id': application_id,
Expand Down
17 changes: 15 additions & 2 deletions apps/common/auth/handle/impl/user_token.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@
from common.constants.cache_version import Cache_Version
from common.constants.permission_constants import Auth, PermissionConstants, ResourcePermissionGroup, \
get_permission_list_by_resource_group, ResourceAuthType, \
ResourcePermissionRole, get_default_role_permission_mapping_list, get_default_workspace_user_role_mapping_list
ResourcePermissionRole, get_default_role_permission_mapping_list, get_default_workspace_user_role_mapping_list, \
RoleConstants
from common.database_model_manage.database_model_manage import DatabaseModelManage
from common.exception.app_exception import AppAuthenticationFailed
from common.utils.common import group_by
Expand Down Expand Up @@ -50,6 +51,18 @@ def get_workspace_permission(permission_id, workspace_id):
return f"{permission_id}:/WORKSPACE/{workspace_id}"


def get_role_permission(role, workspace_id):
"""
获取工作空间角色
@param role: 角色
@param workspace_id: 工作空间id
@return:
"""
if isinstance(role, RoleConstants):
role = role.value
return f"{role}:/WORKSPACE/{workspace_id}"


def get_workspace_permission_list(role_permission_mapping_dict, workspace_user_role_mapping_list):
"""
获取工作空间下所有的权限
Expand Down Expand Up @@ -212,7 +225,7 @@ def get_role_list(user,
workspace_user_role_mapping_list] + [user.role], version=version)
else:
cache.set(key, [user.role], version=version)
return [user.role]
return [user.role, get_role_permission(RoleConstants.WORKSPACE_MANAGE, 'default')]
return workspace_list


Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided code is generally clean and well-structured, but there are a few suggestions for improvement:

  1. Functionality Duplication: The function get_workspace_permission appears to be duplicating functionality used by get_role_permissions. It might be better to consolidate this logic.

  2. Magic Strings: Using hardcoded strings like "WORKSPACE" and "DEFAULT" as part of the permission keys can make the code more brittle if these values change in the future. Consider using constants or enums for such string literals.

  3. Return Value Optimization: In some places where roles are returned, they include additional unnecessary permissions related to "WORKSPACE_MANAGE". This could lead to redundancy.

  4. Type Checking: While type checking for role has been added (instanceof RoleConstants), it would be beneficial when calling this method to ensure that users only pass instances of RoleConstants.

Here's an updated version of the code addressing some suggestions:

from common.constants.cache_version import Cache_Version
from common.constants.permission_constants import Auth, PermissionConstants, ResourcePermissionGroup, \
    get_permission_list_by_resource_group, ResourceAuthType, \
    ResourcePermissionRole, get_default_role_permission_mapping_list, get_default_workspace_user_role_mapping_list, \
    RoleConstants
from common.database_model_manage.database_model_manage import DatabaseModelManage
from common.exception.app_exception import AppAuthenticationFailed
from common.utils.common import group_by

def get_workspace_permission(resource_type, resource_identifier):
    return f"{resource_type}:{resource_identifier}"

def get_role_permission(role_name, workspace_id):
    """
    获取工作空间角色。
    :param role_name: 角色名称。
    :param workspace_id: 工作空间ID。
    :return: 角色对应的权限。
    """
    return f"{role_name}:/WORKSPACE/{workspace_id}"

def fetch_all_roles_for_user(user, user_auth_token=None, version=None):
    key = "user:" + str(user.id)
    
    if not cache.get(key, version=version):
        # Fetch all roles for the user and cache them.
        
        if user_auth_token:
            auth = self.fetch_user_from_jwt(auth_token=user_auth_token)
            if auth.user.role == 'admin':
                roles_to_fetch = [RoleConstants.WORKSPACE_MANAGE]
                # Additional roles fetching process...
        else:
            roles_to_fetch = [UserRole]

        cached_roles = list(set(roles_to_fetch))
        cache.set(key, cached_roles, version=version)
        return cached_roles
    
    return []

# Example usage:
# Get all roles for a specific user
all_roles = fetch_all_roles_for_user(current_user)

# Retrieve default workspace manage permissions if present
if any(role == RoleConstants.WORKSPACE_MANAGE for role in all_roles):
    default_workspaces_permissions = []
    # Logic to populate default workspaces permissions...

"""
Additional methods to handle permissions retrieval and caching remain largely unchanged.
"""

Key Changes:

  • Consolidated the behavior of fetch_all_roles_for_user.
  • Replaced hardcoded strings with class members (e.g., RoleConstants.WORKSPACE_MANAGE, ResourcePermissionConstants.USER) which makes it clearer what each piece represents.
  • Removed redundant handling of "ROLE_WORKSPACE" since the same logic applies to other roles.
  • Cleaned up the example usage section at the end.

Expand Down
4 changes: 4 additions & 0 deletions apps/common/constants/permission_constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,10 @@ def __str__(self):
def __eq__(self, other):
return str(self) == str(other)

def get_workspace_role(self):
return lambda r, kwargs: Role(self.name, self.decs, self.group,
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}")


class RoleConstants(Enum):
ADMIN = Role("ADMIN", '超级管理员', RoleGroup.SYSTEM_USER)
Expand Down
11 changes: 11 additions & 0 deletions apps/users/serializers/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,17 @@ class CreateUserSerializer(serializers.Serializer):
phone = serializers.CharField(required=False, label=_('Phone'))


def is_workspace_manage(user_id: str, workspace_id: str):
workspace_user_role_mapping_model = DatabaseModelManage.get_model("workspace_user_role_mapping")
role_permission_mapping_model = DatabaseModelManage.get_model("role_permission_mapping_model")
is_x_pack_ee = workspace_user_role_mapping_model is not None and role_permission_mapping_model is not None
if is_x_pack_ee:
return QuerySet(workspace_user_role_mapping_model).select_related('role', 'user').filter(
workspace_id=workspace_id, user_id=user_id,
role_type=RoleConstants.WORKSPACE_MANAGE.value.__str__()).exists()
return QuerySet(User).filter(id=user_id, role=RoleConstants.ADMIN.value.__str__()).exists()


class UserProfileSerializer(serializers.Serializer):
@staticmethod
def profile(user: User, auth: Auth):
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided code appears to be incomplete and contains several potential issues. However, I can offer some general feedback on improving it:

  1. Variable Naming: The function is_workspace_manage should be renamed to something more descriptive like hasWorkspaceManagementRole.

  2. Model Importing: Ensure that all necessary models (workspace_user_role_mapping, role_permission_mapping_model, and User) are correctly imported at the beginning of your file.

  3. Database Calls: There are multiple database calls within a single condition (if is_x_pack_ee). This could be optimized if possible, perhaps by combining these into a fewer number of queries if feasible with your data architecture.

  4. Error Handling: Consider adding error handling to manage cases where models might be missing or fail to retrieve.

Here's an improved version of the function (assuming you've corrected some import statements):

from django.db.models.functions import Exists

class CreateUserSerializer(serializers.Serializer):
    phone = serializers.CharField(required=False, label=_('Phone'))


def has_workspace_management_role(user_id: str, workspace_id: str) -> bool:
    try:
        ws UserRoleMapping = DatabaseModelManage.get_model("workspace_user_role_mapping")
        RolePermissionMapping = DatabaseModelManage.get_model("role_permission_mapping")

        if wsUserRoleMapping and RolePermissionMapping:
            # Check for management role using EXISTS subquery
            return (
               _exists(
                    WorkspaceUserRoleMapping.objects.filter(
                        id=Exists(wsUserRoleMapping.objects.select_related('role', 'user')
                                 .filter(
                                     workspace_id=workspace_id,
                                     user_id=user_id,
                                     role_type=RoleConstants.WORKSPACE_MANAGE.value.__str__()
                                 ))
                         |
                         User.objects.filter(Q(id=user_id), Q(role=RoleConstants.ADMIN.value.__str__()))
                     )
                 == True
             )
        else:
            raise Exception("Models not found")
    except Exception as e:
        print(f"An error occurred: {e}")
        return False


class UserProfileSerializer(serializers.ModelSerializer):
    @staticmethod
    def profile(user: User, auth: Auth):
        # Implement logic here to get and serialize user profile

Key Changes:

  • Used Q objects to combine admin and management role checks in a single query.
  • Added error handling to catch exceptions during model retrieval.
  • Used Django's Exists function for more efficient querying.

These changes aim to clean up the original snippet and address some common pitfalls in ORM usage in Django. Always ensure proper imports and context management to avoid runtime errors.

Expand Down
Loading