From 5055a9f352e7d670b49aa16a7eee5703c6646de5 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Tue, 24 Sep 2019 15:18:12 +0800
Subject: [PATCH] =?UTF-8?q?[Update]=20=E4=BF=AE=E6=94=B9=E8=8E=B7=E5=8F=96?=
 =?UTF-8?q?=E8=B5=84=E4=BA=A7=E7=9A=84=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/assets/models/node.py            | 10 ++++++++--
 apps/assets/models/user.py            | 16 +++++-----------
 apps/perms/models/asset_permission.py | 15 +++++----------
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/apps/assets/models/node.py b/apps/assets/models/node.py
index 5d2148314..46e06c63e 100644
--- a/apps/assets/models/node.py
+++ b/apps/assets/models/node.py
@@ -302,13 +302,19 @@ class NodeAssetsMixin:
         return Asset.objects.filter(nodes__key__regex=pattern).distinct()
 
     @classmethod
-    def get_nodes_all_assets(cls, nodes_keys, extra_assets_ids=None):
-        from .asset import Asset
+    def get_nodes_all_assets_ids(cls, nodes_keys):
         nodes_keys = cls.clean_children_keys(nodes_keys)
         assets_ids = set()
         for key in nodes_keys:
             node_assets_ids = cls.tree().all_assets(key)
             assets_ids.update(set(node_assets_ids))
+        return assets_ids
+
+    @classmethod
+    def get_nodes_all_assets(cls, nodes_keys, extra_assets_ids=None):
+        from .asset import Asset
+        nodes_keys = cls.clean_children_keys(nodes_keys)
+        assets_ids = cls.get_nodes_all_assets_ids(nodes_keys)
         if extra_assets_ids:
             assets_ids.update(set(extra_assets_ids))
         return Asset.objects.filter(id__in=assets_ids)
diff --git a/apps/assets/models/user.py b/apps/assets/models/user.py
index 361d770c5..d3447b8ba 100644
--- a/apps/assets/models/user.py
+++ b/apps/assets/models/user.py
@@ -148,18 +148,12 @@ class SystemUser(AssetUser):
         return True, None
 
     def get_all_assets(self):
-        from .node import Node
-        args = [Q(systemuser=self)]
-        pattern = set()
+        from assets.models import Node
         nodes_keys = self.nodes.all().values_list('key', flat=True)
-        nodes_keys = Node.clean_children_keys(nodes_keys)
-        for key in nodes_keys:
-            pattern.add(r'^{0}$|^{0}:'.format(key))
-        pattern = '|'.join(list(pattern))
-        if pattern:
-            args.append(Q(nodes__key__regex=pattern))
-        args = reduce(lambda x, y: x | y, args)
-        assets = Asset.objects.filter(args).distinct()
+        assets_ids = set(self.assets.all().values_list('id', flat=True))
+        nodes_assets_ids = Node.get_nodes_all_assets_ids(nodes_keys)
+        assets_ids.update(nodes_assets_ids)
+        assets = Asset.objects.filter(id__in=assets_ids)
         return assets
 
     class Meta:
diff --git a/apps/perms/models/asset_permission.py b/apps/perms/models/asset_permission.py
index d724d148b..611e6d0de 100644
--- a/apps/perms/models/asset_permission.py
+++ b/apps/perms/models/asset_permission.py
@@ -94,15 +94,10 @@ class AssetPermission(BasePermission):
         )
 
     def get_all_assets(self):
-        args = [Q(granted_by_permissions=self)]
-        pattern = set()
+        from assets.models import Node
         nodes_keys = self.nodes.all().values_list('key', flat=True)
-        nodes_keys = Node.clean_children_keys(nodes_keys)
-        for key in nodes_keys:
-            pattern.add(r'^{0}$|^{0}:'.format(key))
-        pattern = '|'.join(list(pattern))
-        if pattern:
-            args.append(Q(nodes__key__regex=pattern))
-        args = reduce(lambda x, y: x | y, args)
-        assets = Asset.objects.filter(args).distinct()
+        assets_ids = set(self.assets.all().values_list('id', flat=True))
+        nodes_assets_ids = Node.get_nodes_all_assets_ids(nodes_keys)
+        assets_ids.update(nodes_assets_ids)
+        assets = Asset.objects.filter(id__in=assets_ids)
         return assets