refactor: 重构重建用户授权树工具 (#9219)

* perf: 优化 <UserGrantedTreeBuildUtils> 用户授权树构建工具

* feat: 完成计算授权节点资产数量

* refactor: 重构重建用户授权树工具

* merge: v3

Co-authored-by: Bai <baijiangjie@gmail.com>

apps/perms/models/asset_permission.py

@@ -15,6 +15,7 @@ from common.db.models import UnionQuerySet
 from common.utils import date_expired_default
 
 from perms.const import ActionChoices
+from .perm_node import PermNode
 
 __all__ = ['AssetPermission', 'ActionChoices']
 
@@ -48,6 +49,10 @@ class AssetPermissionManager(OrgManager):
     def valid(self):
         return self.get_queryset().valid()
 
+    def get_expired_permissions(self):
+        now = local_now()
+        return self.get_queryset().filter(Q(date_start__lte=now) | Q(date_expired__gte=now))
+
 
 class AssetPermission(OrgModelMixin):
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
@@ -147,10 +152,3 @@ class AssetPermission(OrgModelMixin):
         if flat:
             return user_ids
         return User.objects.filter(id__in=user_ids)
-
-    @classmethod
-    def get_expired_permissions(cls):
-        now = local_now()
-        return cls.objects.filter(Q(date_start__lte=now) | Q(date_expired__gte=now))
-
-

apps/perms/models/perm_node.py

@@ -3,17 +3,20 @@ from django.utils.translation import ugettext_lazy as _
 from django.db import models
 from django.db.models import F, TextChoices
 
-from common.utils import lazyproperty
-from common.db.models import BaseCreateUpdateModel
 from assets.models import Asset, Node, FamilyMixin, Account
 from orgs.mixins.models import OrgModelMixin
+from common.utils import lazyproperty
+from common.db.models import BaseCreateUpdateModel
+
+
+class NodeFrom(TextChoices):
+    granted = 'granted', 'Direct node granted'
+    child = 'child', 'Have children node'
+    asset = 'asset', 'Direct asset granted'
 
 
 class UserAssetGrantedTreeNodeRelation(OrgModelMixin, FamilyMixin, BaseCreateUpdateModel):
-    class NodeFrom(TextChoices):
+    NodeFrom = NodeFrom
-        granted = 'granted', 'Direct node granted'
-        child = 'child', 'Have children node'
-        asset = 'asset', 'Direct asset granted'
 
     user = models.ForeignKey('users.User', db_constraint=False, on_delete=models.CASCADE)
     node = models.ForeignKey('assets.Node', default=None, on_delete=models.CASCADE,
@@ -46,6 +49,8 @@ class UserAssetGrantedTreeNodeRelation(OrgModelMixin, FamilyMixin, BaseCreateUpd
 
 
 class PermNode(Node):
+    NodeFrom = NodeFrom
+
     class Meta:
         proxy = True
         ordering = []

apps/perms/tasks.py

@@ -29,7 +29,7 @@ logger = get_logger(__file__)
 @tmp_to_root_org()
 def check_asset_permission_expired():
     """ 这里的任务要足够短，不要影响周期任务 """
-    perms = AssetPermission.get_expired_permissions()
+    perms = AssetPermission.objects.get_expired_permissions()
     perm_ids = list(perms.distinct().values_list('id', flat=True))
     logger.info(f'Checking expired permissions: {perm_ids}')
     UserPermTreeExpireUtil().expire_perm_tree_for_perms(perm_ids)

apps/perms/utils/permission.py

@@ -1,6 +1,5 @@
-import django
+from django.db.models import QuerySet
-from django.db.models import QuerySet, Model
+
-from collections.abc import Iterable
 from assets.models import Node, Asset
 from common.utils import get_logger
 
@@ -90,11 +89,6 @@ class AssetPermissionUtil(object):
         perms = self.get_permissions(ids=perm_ids)
         return perms
 
-    @staticmethod
-    def get_permissions(ids):
-        perms = AssetPermission.objects.filter(id__in=ids).order_by('-date_expired')
-        return perms
-
     @staticmethod
     def convert_to_queryset_if_need(objs_or_ids, model):
         if not objs_or_ids:
@@ -107,5 +101,7 @@ class AssetPermissionUtil(object):
         ]
         return model.objects.filter(id__in=ids)
 
-
+    @staticmethod
-
+    def get_permissions(ids):
+        perms = AssetPermission.objects.filter(id__in=ids).order_by('-date_expired')
+        return perms

apps/perms/utils/user_perm_tree.py

@@ -1,28 +1,38 @@
 import time
 from collections import defaultdict
 
+from django.conf import settings
 from django.core.cache import cache
 
-from common.decorator import on_transaction_commit
+from users.models import User
-from common.utils import get_logger
+from assets.models import Asset
-from common.utils.common import lazyproperty, timeit
+from assets.utils import NodeAssetsUtil
 from orgs.models import Organization
 from orgs.utils import (
+    current_org,
     tmp_to_org,
     tmp_to_root_org
 )
+from common.decorator import on_transaction_commit
+from common.utils import get_logger
+from common.utils.common import lazyproperty, timeit
+from common.db.models import output_as_string
+
 from perms.locks import UserGrantedTreeRebuildLock
 from perms.models import (
     AssetPermission,
-    UserAssetGrantedTreeNodeRelation
+    UserAssetGrantedTreeNodeRelation,
+    PermNode
 )
-from perms.utils.user_permission import UserGrantedTreeBuildUtils
+
-from users.models import User
 from .permission import AssetPermissionUtil
 
 logger = get_logger(__name__)
 
-__all__ = ['UserPermTreeRefreshUtil', 'UserPermTreeExpireUtil']
+__all__ = [
+    'UserPermTreeRefreshUtil',
+    'UserPermTreeExpireUtil'
+]
 
 
 class _UserPermTreeCacheMixin:
@@ -51,34 +61,34 @@ class UserPermTreeRefreshUtil(_UserPermTreeCacheMixin):
 
     @timeit
     def refresh_if_need(self, force=False):
-        self.clean_user_perm_tree_nodes_for_legacy_org()
+        self._clean_user_perm_tree_for_legacy_org()
-        to_refresh_orgs = self.orgs if force else self.get_user_need_refresh_orgs()
+        to_refresh_orgs = self.orgs if force else self._get_user_need_refresh_orgs()
         if not to_refresh_orgs:
             logger.info('Not have to refresh orgs')
             return
         with UserGrantedTreeRebuildLock(self.user.id):
             for org in to_refresh_orgs:
-                self.rebuild_user_perm_tree_for_org(org)
+                self._rebuild_user_perm_tree_for_org(org)
-        self.mark_user_orgs_refresh_finished([str(org.id) for org in to_refresh_orgs])
+        self._mark_user_orgs_refresh_finished(to_refresh_orgs)
 
-    def rebuild_user_perm_tree_for_org(self, org):
+    def _rebuild_user_perm_tree_for_org(self, org):
         with tmp_to_org(org):
             start = time.time()
-            UserGrantedTreeBuildUtils(self.user).rebuild_user_granted_tree()
+            UserPermTreeBuildUtil(self.user).rebuild_user_perm_tree()
             end = time.time()
             logger.info(
                 'Refresh user [{user}] org [{org}] perm tree, user {use_time:.2f}s'
-                ''.format(user=self.user, org=org, use_time=end - start)
+                ''.format(user=self.user, org=org, use_time=end-start)
             )
 
-    def clean_user_perm_tree_nodes_for_legacy_org(self):
+    def _clean_user_perm_tree_for_legacy_org(self):
         with tmp_to_root_org():
             """ Clean user legacy org node relations """
             user_relations = UserAssetGrantedTreeNodeRelation.objects.filter(user=self.user)
             user_legacy_org_relations = user_relations.exclude(org_id__in=self.org_ids)
             user_legacy_org_relations.delete()
 
-    def get_user_need_refresh_orgs(self):
+    def _get_user_need_refresh_orgs(self):
         cached_org_ids = self.client.smembers(self.cache_key_user)
         cached_org_ids = {oid.decode() for oid in cached_org_ids}
         to_refresh_org_ids = set(self.org_ids) - cached_org_ids
@@ -86,7 +96,7 @@ class UserPermTreeRefreshUtil(_UserPermTreeCacheMixin):
         logger.info(f'Need to refresh orgs: {to_refresh_orgs}')
         return to_refresh_orgs
 
-    def mark_user_orgs_refresh_finished(self, org_ids):
+    def _mark_user_orgs_refresh_finished(self, org_ids):
         self.client.sadd(self.cache_key_user, *org_ids)
 
 
@@ -141,3 +151,181 @@ class UserPermTreeExpireUtil(_UserPermTreeCacheMixin):
                 p.delete(k)
             p.execute()
         logger.info('Expire all user perm tree')
+
+
+class UserPermTreeBuildUtil(object):
+    node_only_fields = ('id', 'key', 'parent_key', 'org_id')
+
+    def __init__(self, user):
+        self.user = user
+        self.user_perm_ids = AssetPermissionUtil().get_permissions_for_user(self.user, flat=True)
+        # {key: node}
+        self._perm_nodes_key_node_mapper = {}
+
+    def rebuild_user_perm_tree(self):
+        self.clean_user_perm_tree()
+        if not self.user_perm_ids:
+            logger.info('User({}) not have permissions'.format(self.user))
+            return
+        self.compute_perm_nodes()
+        self.compute_perm_nodes_asset_amount()
+        self.create_mapping_nodes()
+
+    def clean_user_perm_tree(self):
+        UserAssetGrantedTreeNodeRelation.objects.filter(user=self.user).delete()
+
+    def compute_perm_nodes(self):
+        self._compute_perm_nodes_for_direct()
+        self._compute_perm_nodes_for_direct_asset_if_need()
+        self._compute_perm_nodes_for_ancestor()
+
+    def compute_perm_nodes_asset_amount(self):
+        """ 这里计算的是一个组织的授权树 """
+        computed = self._only_compute_root_node_assets_amount_if_need()
+        if computed:
+            return
+
+        nodekey_assetid_mapper = defaultdict(set)
+        org_id = current_org.id
+        for key in self.perm_node_keys_for_granted:
+            asset_ids = PermNode.get_all_asset_ids_by_node_key(org_id, key)
+            nodekey_assetid_mapper[key].update(asset_ids)
+
+        for asset_id, node_id in self.direct_asset_id_node_id_pairs:
+            node_key = self.perm_nodes_id_key_mapper.get(node_id)
+            if not node_key:
+                continue
+            nodekey_assetid_mapper[node_key].add(asset_id)
+
+        util = NodeAssetsUtil(self.perm_nodes, nodekey_assetid_mapper)
+        util.generate()
+
+        for node in self.perm_nodes:
+            assets_amount = util.get_assets_amount(node.key)
+            node.assets_amount = assets_amount
+
+    def create_mapping_nodes(self):
+        to_create = []
+        for node in self.perm_nodes:
+            relation = UserAssetGrantedTreeNodeRelation(
+                user=self.user,
+                node=node,
+                node_key=node.key,
+                node_parent_key=node.parent_key,
+                node_from=node.node_from,
+                node_assets_amount=node.assets_amount,
+                org_id=node.org_id
+            )
+            to_create.append(relation)
+
+        UserAssetGrantedTreeNodeRelation.objects.bulk_create(to_create)
+
+    def _compute_perm_nodes_for_direct(self):
+        """ 直接授权的节点（叶子节点）"""
+        for node in self.direct_nodes:
+            if self.has_any_ancestor_direct_permed(node):
+                continue
+            node.node_from = node.NodeFrom.granted
+            self._perm_nodes_key_node_mapper[node.key] = node
+
+    def _compute_perm_nodes_for_direct_asset_if_need(self):
+        """ 直接授权的资产所在的节点（叶子节点）"""
+        if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
+            return
+        for node in self.direct_asset_nodes:
+            if self.has_any_ancestor_direct_permed(node):
+                continue
+            if node.key in self._perm_nodes_key_node_mapper:
+                continue
+            node.node_from = node.NodeFrom.asset
+            self._perm_nodes_key_node_mapper[node.key] = node
+
+    def _compute_perm_nodes_for_ancestor(self):
+        """ 直接授权节点 和 直接授权资产所在节点 的所有祖先节点 (构造完整树) """
+        ancestor_keys = set()
+        for node in self._perm_nodes_key_node_mapper.values():
+            ancestor_keys.update(node.get_ancestor_keys())
+        ancestor_keys -= set(self._perm_nodes_key_node_mapper.keys())
+
+        ancestors = PermNode.objects.filter(key__in=ancestor_keys).only(*self.node_only_fields)
+        for node in ancestors:
+            node.node_from = node.NodeFrom.child
+            self._perm_nodes_key_node_mapper[node.key] = node
+
+    @lazyproperty
+    def perm_node_keys_for_granted(self):
+        keys = [
+            key for key, node in self._perm_nodes_key_node_mapper.items()
+            if node.node_from == node.NodeFrom.granted
+        ]
+        return keys
+
+    @lazyproperty
+    def perm_nodes_id_key_mapper(self):
+        mapper = {
+            node.id.hex: node.key
+            for key, node in self._perm_nodes_key_node_mapper.items()
+        }
+        return mapper
+
+    def _only_compute_root_node_assets_amount_if_need(self):
+        if len(self.perm_nodes) != 1:
+            return False
+        root_node = self.perm_nodes[0]
+        if not root_node.is_org_root():
+            return False
+        if root_node.node_from != root_node.NodeFrom.granted:
+            return False
+        root_node.granted_assets_amount = len(root_node.get_all_asset_ids())
+        return True
+
+    @lazyproperty
+    def perm_nodes(self):
+        """ 授权树的所有节点 """
+        return list(self._perm_nodes_key_node_mapper.values())
+
+    def has_any_ancestor_direct_permed(self, node):
+        """ 任何一个祖先节点被直接授权 """
+        return bool(set(node.get_ancestor_keys()) & set(self.direct_node_keys))
+
+    @lazyproperty
+    def direct_node_keys(self):
+        """ 直接授权的节点 keys """
+        return {n.key for n in self.direct_nodes}
+
+    @lazyproperty
+    def direct_nodes(self):
+        """ 直接授权的节点 """
+        node_ids = AssetPermission.nodes.through.objects \
+            .filter(assetpermission_id__in=self.user_perm_ids) \
+            .values_list('node_id', flat=True).distinct()
+        nodes = PermNode.objects.filter(id__in=node_ids).only(*self.node_only_fields)
+        return nodes
+
+    @lazyproperty
+    def direct_asset_nodes(self):
+        """ 获取直接授权的资产所在的节点 """
+        node_ids = [node_id for asset_id, node_id in self.direct_asset_id_node_id_pairs]
+        nodes = PermNode.objects.filter(id__in=node_ids).distinct().only(*self.node_only_fields)
+        return nodes
+
+    @lazyproperty
+    def direct_asset_id_node_id_pairs(self):
+        """ 直接授权的资产 id 和 节点 id  """
+        asset_node_pairs = Asset.nodes.through.objects \
+            .filter(asset_id__in=self.direct_asset_ids) \
+            .annotate(
+                str_asset_id=output_as_string('asset_id'),
+                str_node_id=output_as_string('node_id')
+            ).values_list('str_asset_id', 'str_node_id')
+        asset_node_pairs = list(asset_node_pairs)
+        return asset_node_pairs
+
+    @lazyproperty
+    def direct_asset_ids(self):
+        """ 直接授权的资产 ids """
+        asset_ids = AssetPermission.assets.through.objects \
+            .filter(assetpermission_id__in=self.user_perm_ids) \
+            .values_list('asset_id', flat=True) \
+            .distinct()
+        return asset_ids

apps/perms/utils/user_permission.py

@@ -51,217 +51,6 @@ class UserGrantedUtilsBase:
         return asset_perm_ids
 
 
-class UserGrantedTreeBuildUtils(UserGrantedUtilsBase):
-
-    def get_direct_granted_nodes(self) -> NodeQuerySet:
-        # 查询直接授权节点
-        nodes = PermNode.objects.filter(
-            granted_by_permissions__id__in=self.asset_perm_ids
-        ).distinct()
-        return nodes
-
-    @lazyproperty
-    def direct_granted_asset_ids(self) -> list:
-        # 3.15
-        asset_ids = AssetPermission.assets.through.objects.filter(
-            assetpermission_id__in=self.asset_perm_ids
-        ).annotate(
-            asset_id_str=output_as_string('asset_id')
-        ).values_list(
-            'asset_id_str', flat=True
-        ).distinct()
-
-        asset_ids = list(asset_ids)
-        return asset_ids
-
-    @ensure_in_real_or_default_org
-    def rebuild_user_granted_tree(self):
-        """
-        注意：调用该方法一定要被 `UserGrantedTreeRebuildLock` 锁住
-        """
-        user = self.user
-
-        # 先删除旧的授权树🌲
-        UserAssetGrantedTreeNodeRelation.objects.filter(user=user).delete()
-
-        if not self.asset_perm_ids:
-            # 没有授权直接返回
-            return
-
-        nodes = self.compute_perm_nodes_tree()
-        self.compute_node_assets_amount(nodes)
-        if not nodes:
-            return
-        self.create_mapping_nodes(nodes)
-
-    @timeit
-    def compute_perm_nodes_tree(self, node_only_fields=NODE_ONLY_FIELDS) -> list:
-
-        # 查询直接授权节点
-        nodes = self.get_direct_granted_nodes().only(*node_only_fields)
-        nodes = list(nodes)
-
-        # 授权的节点 key 集合
-        granted_key_set = {_node.key for _node in nodes}
-
-        def _has_ancestor_granted(node: PermNode):
-            """
-            判断一个节点是否有授权过的祖先节点
-            """
-            ancestor_keys = set(node.get_ancestor_keys())
-            return ancestor_keys & granted_key_set
-
-        key2leaf_nodes_mapper = {}
-
-        # 给授权节点设置 granted 标识，同时去重
-        for node in nodes:
-            node: PermNode
-            if _has_ancestor_granted(node):
-                continue
-            node.node_from = NodeFrom.granted
-            key2leaf_nodes_mapper[node.key] = node
-
-        # 查询授权资产关联的节点设置
-        def process_direct_granted_assets():
-            # 查询直接授权资产
-            node_ids = {node_id_str for node_id_str, _ in self.direct_granted_asset_id_node_id_str_pairs}
-            # 查询授权资产关联的节点设置 2.80
-            granted_asset_nodes = PermNode.objects.filter(
-                id__in=node_ids
-            ).distinct().only(*node_only_fields)
-            granted_asset_nodes = list(granted_asset_nodes)
-
-            # 给资产授权关联的节点设置 is_asset_granted 标识，同时去重
-            for node in granted_asset_nodes:
-                if _has_ancestor_granted(node):
-                    continue
-                if node.key in key2leaf_nodes_mapper:
-                    continue
-                node.node_from = NodeFrom.asset
-                key2leaf_nodes_mapper[node.key] = node
-
-        if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
-            process_direct_granted_assets()
-
-        leaf_nodes = key2leaf_nodes_mapper.values()
-
-        # 计算所有祖先节点
-        ancestor_keys = set()
-        for node in leaf_nodes:
-            ancestor_keys.update(node.get_ancestor_keys())
-
-        # 从祖先节点 key 中去掉同时也是叶子节点的 key
-        ancestor_keys -= key2leaf_nodes_mapper.keys()
-        # 查出祖先节点
-        ancestors = PermNode.objects.filter(key__in=ancestor_keys).only(*node_only_fields)
-        ancestors = list(ancestors)
-        for node in ancestors:
-            node.node_from = NodeFrom.child
-        result = [*leaf_nodes, *ancestors]
-        return result
-
-    @timeit
-    def create_mapping_nodes(self, nodes):
-        user = self.user
-        to_create = []
-
-        for node in nodes:
-            to_create.append(UserAssetGrantedTreeNodeRelation(
-                user=user,
-                node=node,
-                node_key=node.key,
-                node_parent_key=node.parent_key,
-                node_from=node.node_from,
-                node_assets_amount=node.assets_amount,
-                org_id=node.org_id
-            ))
-
-        UserAssetGrantedTreeNodeRelation.objects.bulk_create(to_create)
-
-    @timeit
-    def _fill_direct_granted_node_asset_ids_from_mem(self, nodes_key, mapper):
-        org_id = current_org.id
-        for key in nodes_key:
-            asset_ids = PermNode.get_all_asset_ids_by_node_key(org_id, key)
-            mapper[key].update(asset_ids)
-
-    @lazyproperty
-    def direct_granted_asset_id_node_id_str_pairs(self):
-        node_asset_pairs = Asset.nodes.through.objects.filter(
-            asset_id__in=self.direct_granted_asset_ids
-        ).annotate(
-            asset_id_str=output_as_string('asset_id'),
-            node_id_str=output_as_string('node_id')
-        ).values_list(
-            'node_id_str', 'asset_id_str'
-        )
-        node_asset_pairs = list(node_asset_pairs)
-        return node_asset_pairs
-
-    @timeit
-    def compute_node_assets_amount(self, nodes: List[PermNode]):
-        """
-        这里计算的是一个组织的
-        """
-        # 直接授权了根节点，直接计算
-        if len(nodes) == 1:
-            node = nodes[0]
-            if node.node_from == NodeFrom.granted and node.key.isdigit():
-                with tmp_to_org(node.org):
-                    node.granted_assets_amount = len(node.get_all_asset_ids())
-                    return
-
-        direct_granted_nodes_key = []
-        node_id_key_mapper = {}
-        for node in nodes:
-            if node.node_from == NodeFrom.granted:
-                direct_granted_nodes_key.append(node.key)
-            node_id_key_mapper[node.id.hex] = node.key
-
-        # 授权的节点和直接资产的映射
-        nodekey_assetsid_mapper = defaultdict(set)
-        # 直接授权的节点，资产从完整树过来
-        self._fill_direct_granted_node_asset_ids_from_mem(
-            direct_granted_nodes_key, nodekey_assetsid_mapper
-        )
-
-        # 处理直接授权资产
-        # 直接授权资产，取节点与资产的关系
-        node_asset_pairs = self.direct_granted_asset_id_node_id_str_pairs
-        node_asset_pairs = list(node_asset_pairs)
-
-        for node_id, asset_id in node_asset_pairs:
-            if node_id not in node_id_key_mapper:
-                continue
-            node_key = node_id_key_mapper[node_id]
-            nodekey_assetsid_mapper[node_key].add(asset_id)
-
-        util = NodeAssetsUtil(nodes, nodekey_assetsid_mapper)
-        util.generate()
-
-        for node in nodes:
-            assets_amount = util.get_assets_amount(node.key)
-            node.assets_amount = assets_amount
-
-    def get_whole_tree_nodes(self) -> list:
-        node_only_fields = NODE_ONLY_FIELDS + ('value', 'full_value')
-        nodes = self.compute_perm_nodes_tree(node_only_fields=node_only_fields)
-        self.compute_node_assets_amount(nodes)
-
-        # 查询直接授权节点的子节点
-        q = Q()
-        for node in self.get_direct_granted_nodes().only('key'):
-            q |= Q(key__startswith=f'{node.key}:')
-
-        if q:
-            descendant_nodes = PermNode.objects.filter(q).distinct()
-        else:
-            descendant_nodes = PermNode.objects.none()
-
-        nodes.extend(descendant_nodes)
-        return nodes
-
-
 class UserGrantedAssetsQueryUtils(UserGrantedUtilsBase):
 
     def get_favorite_assets(self) -> QuerySet:

apps/users/models/user.py

@@ -742,13 +742,6 @@ class User(AuthMixin, TokenMixin, RoleMixin, MFAMixin, AbstractUser):
     def __str__(self):
         return '{0.name}({0.username})'.format(self)
 
-    @classmethod
-    def get_group_ids_by_user_id(cls, user_id):
-        group_ids = cls.groups.through.objects.filter(user_id=user_id) \
-            .distinct().values_list('usergroup_id', flat=True)
-        group_ids = list(group_ids)
-        return group_ids
-
     @property
     def receive_backends(self):
         return self.user_msg_subscription.receive_backends