perf(User): 用户列表在大规模数据情况下慢

Co-authored-by: xinwen <coderWen@126.com>

apps/assets/api/node.py

@@ -201,10 +201,8 @@ class NodeAddChildrenApi(generics.UpdateAPIView):
     def put(self, request, *args, **kwargs):
         instance = self.get_object()
         nodes_id = request.data.get("nodes")
-        children = [get_object_or_none(Node, id=pk) for pk in nodes_id]
+        children = Node.objects.filter(id__in=nodes_id)
         for node in children:
-            if not node:
-                continue
             node.parent = instance
         return Response("OK")
 

apps/assets/models/node.py

@@ -103,7 +103,7 @@ class FamilyMixin:
             if value is None:
                 value = child_key
             child = self.__class__.objects.create(
-                id=_id, key=child_key, value=value, parent_key=self.key,
+                id=_id, key=child_key, value=value
             )
             return child
 

apps/assets/signals_handler.py

@@ -4,7 +4,7 @@ from operator import add, sub
 
 from assets.utils import is_asset_exists_in_node
 from django.db.models.signals import (
-    post_save, m2m_changed, pre_delete, post_delete
+    post_save, m2m_changed, pre_delete, post_delete, pre_save
 )
 from django.db.models import Q, F
 from django.dispatch import receiver
@@ -37,6 +37,11 @@ def test_asset_conn_on_created(asset):
     test_asset_connectivity_util.delay([asset])
 
 
+@receiver(pre_save, sender=Node)
+def on_node_pre_save(sender, instance: Node, **kwargs):
+    instance.parent_key = instance.compute_parent_key()
+
+
 @receiver(post_save, sender=Asset)
 @on_transaction_commit
 def on_asset_created_or_update(sender, instance=None, created=False, **kwargs):
@@ -91,22 +96,24 @@ def on_system_user_assets_change(instance, action, model, pk_set, **kwargs):
 
 
 @receiver(m2m_changed, sender=SystemUser.users.through)
-def on_system_user_users_change(sender, instance=None, action='', model=None, pk_set=None, **kwargs):
+def on_system_user_users_change(sender, instance: SystemUser, action, model, pk_set, reverse, **kwargs):
     """
     当系统用户和用户关系发生变化时，应该重新推送系统用户资产中
     """
     if action != POST_ADD:
         return
+
+    if reverse:
+        raise M2MReverseNotAllowed
+
     if not instance.username_same_with_user:
         return
+
     logger.debug("System user users change signal recv: {}".format(instance))
-    queryset = model.objects.filter(pk__in=pk_set)
-    if model == SystemUser:
-        system_users = queryset
-    else:
-        system_users = [instance]
-    for s in system_users:
-        push_system_user_to_assets_manual.delay(s)
+    usernames = model.objects.filter(pk__in=pk_set).values_list('username', flat=True)
+
+    for username in usernames:
+        push_system_user_to_assets_manual.delay(instance, username)
 
 
 @receiver(m2m_changed, sender=SystemUser.nodes.through)

apps/assets/tasks/nodes_amount.py

@@ -7,7 +7,7 @@ from common.utils import get_logger
 logger = get_logger(__file__)
 
 
-@register_as_period_task(crontab='* 2 * * *')
+@register_as_period_task(crontab='0 2 * * *')
 @shared_task(queue='celery_heavy_tasks')
 def check_node_assets_amount_celery_task():
     check_node_assets_amount()

apps/orgs/api.py

@@ -92,6 +92,7 @@ class OrgMemberAdminRelationBulkViewSet(JMSBulkRelationModelViewSet):
     serializer_class = OrgMemberAdminSerializer
     filterset_class = OrgMemberRelationFilterSet
     search_fields = ('user__name', 'user__username', 'org__name')
+    lookup_field = 'user_id'
 
     def get_queryset(self):
         queryset = super().get_queryset()
@@ -116,6 +117,7 @@ class OrgMemberUserRelationBulkViewSet(JMSBulkRelationModelViewSet):
     serializer_class = OrgMemberUserSerializer
     filterset_class = OrgMemberRelationFilterSet
     search_fields = ('user__name', 'user__username', 'org__name')
+    lookup_field = 'user_id'
 
     def get_queryset(self):
         queryset = super().get_queryset()

apps/perms/signals_handler.py

@@ -6,7 +6,7 @@ from django.dispatch import receiver
 from perms.tasks import create_rebuild_user_tree_task, \
     create_rebuild_user_tree_task_by_related_nodes_or_assets
 from users.models import User, UserGroup
-from assets.models import Asset
+from assets.models import Asset, SystemUser
 from common.utils import get_logger
 from common.exceptions import M2MReverseNotAllowed
 from common.const.signals import POST_ADD, POST_REMOVE, POST_CLEAR
@@ -16,6 +16,42 @@ from .models import AssetPermission, RemoteAppPermission
 logger = get_logger(__file__)
 
 
+def handle_rebuild_user_tree(instance, action, reverse, pk_set, **kwargs):
+    if action.startswith('post'):
+        if reverse:
+            create_rebuild_user_tree_task(pk_set)
+        else:
+            create_rebuild_user_tree_task([instance.id])
+
+
+def handle_bind_groups_systemuser(instance, action, reverse, pk_set, **kwargs):
+    """
+    UserGroup 增加 User 时，增加的 User 需要与 UserGroup 关联的动态系统用户相关联
+    """
+    user: User
+
+    if action != POST_ADD:
+        return
+
+    if not reverse:
+        # 一个用户添加了多个用户组
+        users_id = [instance.id]
+        system_users = SystemUser.objects.filter(groups__id__in=pk_set).distinct()
+    else:
+        # 一个用户组添加了多个用户
+        users_id = pk_set
+        system_users = SystemUser.objects.filter(groups__id=instance.pk).distinct()
+
+    for system_user in system_users:
+        system_user.users.add(*users_id)
+
+
+@receiver(m2m_changed, sender=User.groups.through)
+def on_user_groups_change(**kwargs):
+    handle_rebuild_user_tree(**kwargs)
+    handle_bind_groups_systemuser(**kwargs)
+
+
 @receiver([pre_save], sender=AssetPermission)
 def on_asset_perm_deactive(instance: AssetPermission, **kwargs):
     try:

apps/users/api/user.py

@@ -6,8 +6,8 @@ from rest_framework.decorators import action
 from rest_framework import generics
 from rest_framework.response import Response
 from rest_framework_bulk import BulkModelViewSet
+from django.db.models import Prefetch
 
-from common.db.aggregates import GroupConcat
 from common.permissions import (
     IsOrgAdmin, IsOrgAdminOrAppUser,
     CanUpdateDeleteUser, IsSuperUser
@@ -44,9 +44,18 @@ class UserViewSet(CommonApiMixin, UserQuerysetMixin, BulkModelViewSet):
     extra_filter_backends = [OrgRoleUserFilterBackend]
 
     def get_queryset(self):
-        return super().get_queryset().annotate(
-            gc_m2m_org_members__role=GroupConcat('m2m_org_members__role'),
-        ).prefetch_related('groups')
+        queryset = super().get_queryset().prefetch_related(
+            'groups'
+        )
+        if current_org.is_real():
+            # 为在列表中计算用户在真实组织里的角色
+            queryset = queryset.prefetch_related(
+                Prefetch(
+                    'm2m_org_members',
+                    queryset=OrganizationMember.objects.filter(org__id=current_org.id)
+                )
+            )
+        return queryset
 
     def send_created_signal(self, users):
         if not isinstance(users, list):

apps/users/models/user.py

@@ -170,22 +170,18 @@ class RoleMixin:
         from orgs.models import ROLE as ORG_ROLE
 
         if not current_org.is_real():
+            # 不是真实的组织，取 User 本身的角色
             if self.is_superuser:
                 return [ORG_ROLE.ADMIN]
             else:
                 return [ORG_ROLE.USER]
 
-        if hasattr(self, 'gc_m2m_org_members__role'):
-            names = self.gc_m2m_org_members__role
-            if isinstance(names, str):
-                roles = set(self.gc_m2m_org_members__role.split(','))
-            else:
-                roles = set()
-        else:
-            roles = set(self.m2m_org_members.filter(
-                org_id=current_org.id
-            ).values_list('role', flat=True))
-        roles = list(roles)
+        # 是真实组织，取 OrganizationMember 中的角色
+        roles = [
+            org_member.role
+            for org_member in self.m2m_org_members.all()
+            if org_member.org_id == current_org.id
+        ]
         roles.sort()
         return roles
 

apps/users/signals_handler.py

@@ -2,14 +2,12 @@
 #
 
 from django.dispatch import receiver
-from django.db.models.signals import m2m_changed
 from django_auth_ldap.backend import populate_user
 from django.conf import settings
 from django_cas_ng.signals import cas_user_authenticated
 
 from jms_oidc_rp.signals import openid_create_or_update_user
 
-from perms.tasks import create_rebuild_user_tree_task
 from common.utils import get_logger
 from .signals import post_user_create
 from .models import User
@@ -27,15 +25,6 @@ def on_user_create(sender, user=None, **kwargs):
         send_user_created_mail(user)
 
 
-@receiver(m2m_changed, sender=User.groups.through)
-def on_user_groups_change(instance, action, reverse, pk_set, **kwargs):
-    if action.startswith('post'):
-        if reverse:
-            create_rebuild_user_tree_task(pk_set)
-        else:
-            create_rebuild_user_tree_task([instance.id])
-
-
 @receiver(cas_user_authenticated)
 def on_cas_user_authenticated(sender, user, created, **kwargs):
     if created: