From 78bf6f5817126a41255b066900774796df5f0e56 Mon Sep 17 00:00:00 2001
From: xinwen <coderWen@126.com>
Date: Wed, 3 Mar 2021 15:36:42 +0800
Subject: [PATCH] =?UTF-8?q?refactor:=20=E8=8E=B7=E5=8F=96=E6=8E=88?=
 =?UTF-8?q?=E6=9D=83=E6=A0=91=E6=88=96=E8=80=85=E8=B5=84=E4=BA=A7=E5=88=97?=
 =?UTF-8?q?=E8=A1=A8=E6=97=B6=E9=81=BF=E5=85=8D=E8=AF=BB=E6=97=B6=E9=94=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/perms/pagination.py                  |  8 +++-
 apps/perms/utils/asset/user_permission.py | 50 ++++++++++++-----------
 2 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/apps/perms/pagination.py b/apps/perms/pagination.py
index c740830c2..248958a3e 100644
--- a/apps/perms/pagination.py
+++ b/apps/perms/pagination.py
@@ -29,8 +29,12 @@ class AllGrantedAssetPagination(GrantedAssetPaginationBase):
     def get_count_from_nodes(self, queryset):
         if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
             return None
-        assets_amount = sum(UserAssetGrantedTreeNodeRelation.objects.filter(
+        values = UserAssetGrantedTreeNodeRelation.objects.filter(
             user=self._user, node_parent_key=''
-        ).values_list('node_assets_amount', flat=True))
+        ).values_list('node_assets_amount', flat=True)
+        if not values:
+            return None
+
+        assets_amount = sum(values)
         logger.debug(f'Hit all assets amount {assets_amount} -> {self._request.get_full_path()}')
         return assets_amount
diff --git a/apps/perms/utils/asset/user_permission.py b/apps/perms/utils/asset/user_permission.py
index 518148a1b..14205e12e 100644
--- a/apps/perms/utils/asset/user_permission.py
+++ b/apps/perms/utils/asset/user_permission.py
@@ -1,5 +1,6 @@
 from collections import defaultdict
 from typing import List, Tuple
+import time
 
 from django.core.cache import cache
 from django.conf import settings
@@ -185,11 +186,16 @@ class UserGrantedTreeRefreshController:
         return {org_id.decode() for org_id in org_ids}
 
     def set_all_orgs_as_builed(self):
-        orgs_id = [str(org_id) for org_id in self.orgs_id]
-        self.client.sadd(self.key, *orgs_id)
+        self.client.sadd(self.key, *self.orgs_id)
+
+    def have_need_refresh_orgs(self):
+        builded_org_ids = self.client.smembers(self.key)
+        builded_org_ids = {org_id.decode() for org_id in builded_org_ids}
+        have = self.orgs_id - builded_org_ids
+        return have
 
     def get_need_refresh_orgs_and_fill_up(self):
-        orgs_id = set(str(org_id) for org_id in self.orgs_id)
+        orgs_id = self.orgs_id
 
         with self.client.pipeline() as p:
             p.smembers(self.key)
@@ -197,8 +203,7 @@ class UserGrantedTreeRefreshController:
             ret = p.execute()
             builded_orgs_id = {org_id.decode() for org_id in ret[0]}
             ids = orgs_id - builded_orgs_id
-            orgs = set()
-            orgs.update(Organization.objects.filter(id__in=ids))
+            orgs = {*Organization.objects.filter(id__in=ids)}
             logger.info(f'Need rebuild orgs are {orgs}, builed orgs are {ret[0]}, all orgs are {orgs_id}')
             return orgs
 
@@ -285,7 +290,7 @@ class UserGrantedTreeRefreshController:
 
     @lazyproperty
     def orgs_id(self):
-        ret = {org.id for org in self.orgs}
+        ret = {str(org.id) for org in self.orgs}
         return ret
 
     @lazyproperty
@@ -297,21 +302,24 @@ class UserGrantedTreeRefreshController:
     def refresh_if_need(self, force=False):
         user = self.user
 
-        with UserGrantedTreeRebuildLock(user_id=user.id):
-            with tmp_to_root_org():
-                UserAssetGrantedTreeNodeRelation.objects.filter(user=user).exclude(org_id__in=self.orgs_id).delete()
-                exists = UserAssetGrantedTreeNodeRelation.objects.filter(user=user).exists()
+        with tmp_to_root_org():
+            UserAssetGrantedTreeNodeRelation.objects.filter(user=user).exclude(org_id__in=self.orgs_id).delete()
 
-            if force or not exists:
-                orgs = self.orgs
-                self.set_all_orgs_as_builed()
-            else:
-                orgs = self.get_need_refresh_orgs_and_fill_up()
+        if force or self.have_need_refresh_orgs():
+            with UserGrantedTreeRebuildLock(user_id=user.id):
+                if force:
+                    orgs = self.orgs
+                    self.set_all_orgs_as_builed()
+                else:
+                    orgs = self.get_need_refresh_orgs_and_fill_up()
 
-            for org in orgs:
-                with tmp_to_org(org):
-                    utils = UserGrantedTreeBuildUtils(user)
-                    utils.rebuild_user_granted_tree()
+                for org in orgs:
+                    with tmp_to_org(org):
+                        t_start = time.time()
+                        logger.info(f'Rebuild user tree: user={self.user} org={current_org}')
+                        utils = UserGrantedTreeBuildUtils(user)
+                        utils.rebuild_user_granted_tree()
+                        logger.info(f'Rebuild user tree ok: cost={time.time() - t_start} user={self.user} org={current_org}')
 
 
 class UserGrantedUtilsBase:
@@ -353,15 +361,11 @@ class UserGrantedTreeBuildUtils(UserGrantedUtilsBase):
         asset_ids = list(asset_ids)
         return asset_ids
 
-    @timeit
     @ensure_in_real_or_default_org
     def rebuild_user_granted_tree(self):
         """
         注意：调用该方法一定要被 `UserGrantedTreeRebuildLock` 锁住
         """
-
-        logger.info(f'Rebuild user tree: user={self.user} org={current_org}')
-
         user = self.user
 
         # 先删除旧的授权树🌲