From 15ac81a422defdf1dbb9e506a712934cdceccfbe Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Wed, 3 Jan 2024 16:14:27 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E6=A0=87=E7=AD=BE?=
 =?UTF-8?q?=E7=BB=91=E5=AE=9A=EF=BC=8C=E4=BB=85=E7=BB=91=E5=AE=9A=E5=88=B0?=
 =?UTF-8?q?=E8=B5=84=E4=BA=A7=E4=B8=8A?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/assets/api/platform.py             |  3 ++-
 apps/assets/serializers/asset/common.py |  9 +++++----
 apps/assets/serializers/platform.py     |  6 ------
 apps/common/drf/filters.py              |  4 ++--
 apps/common/signal_handlers.py          |  2 +-
 apps/labels/api.py                      | 24 ++++++++++++++++--------
 apps/labels/mixins.py                   | 20 ++++++++++++++++++--
 7 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/apps/assets/api/platform.py b/apps/assets/api/platform.py
index 9f3a55bda..6a31f9519 100644
--- a/apps/assets/api/platform.py
+++ b/apps/assets/api/platform.py
@@ -29,8 +29,9 @@ class AssetPlatformViewSet(JMSModelViewSet):
     }
 
     def get_queryset(self):
+        # 因为没有走分页逻辑,所以需要这里 prefetch
         queryset = super().get_queryset().prefetch_related(
-            'protocols', 'automation'
+            'protocols', 'automation', 'labels', 'labels__label',
         )
         queryset = queryset.filter(type__in=AllTypes.get_types_values())
         return queryset
diff --git a/apps/assets/serializers/asset/common.py b/apps/assets/serializers/asset/common.py
index 6c45944de..501d06780 100644
--- a/apps/assets/serializers/asset/common.py
+++ b/apps/assets/serializers/asset/common.py
@@ -100,7 +100,10 @@ class AssetAccountSerializer(AccountSerializer):
     class Meta(AccountSerializer.Meta):
         fields = [
             f for f in AccountSerializer.Meta.fields
-            if f not in ['spec_info']
+            if f not in [
+                'spec_info', 'connectivity', 'labels', 'created_by',
+                'date_update', 'date_created'
+            ]
         ]
         extra_kwargs = {
             **AccountSerializer.Meta.extra_kwargs,
@@ -375,7 +378,6 @@ class AssetSerializer(BulkOrgResourceModelSerializer, ResourceLabelsMixin, Writa
 
 
 class DetailMixin(serializers.Serializer):
-    accounts = AssetAccountSerializer(many=True, required=False, label=_('Accounts'))
     spec_info = MethodSerializer(label=_('Spec info'), read_only=True)
     gathered_info = MethodSerializer(label=_('Gathered info'), read_only=True)
     auto_config = serializers.DictField(read_only=True, label=_('Auto info'))
@@ -390,8 +392,7 @@ class DetailMixin(serializers.Serializer):
     def get_field_names(self, declared_fields, info):
         names = super().get_field_names(declared_fields, info)
         names.extend([
-            'accounts', 'gathered_info', 'spec_info',
-            'auto_config',
+            'gathered_info', 'spec_info', 'auto_config',
         ])
         return names
 
diff --git a/apps/assets/serializers/platform.py b/apps/assets/serializers/platform.py
index f67df9906..2cc7d73ef 100644
--- a/apps/assets/serializers/platform.py
+++ b/apps/assets/serializers/platform.py
@@ -200,12 +200,6 @@ class PlatformSerializer(ResourceLabelsMixin, WritableNestedModelSerializer):
         constraints = AllTypes.get_constraints(category, tp)
         return constraints
 
-    @classmethod
-    def setup_eager_loading(cls, queryset):
-        queryset = queryset.prefetch_related('protocols', 'automation') \
-            .prefetch_related('labels', 'labels__label')
-        return queryset
-
     def validate_protocols(self, protocols):
         if not protocols:
             raise serializers.ValidationError(_("Protocols is required"))
diff --git a/apps/common/drf/filters.py b/apps/common/drf/filters.py
index 0fb8ca63a..d187cd414 100644
--- a/apps/common/drf/filters.py
+++ b/apps/common/drf/filters.py
@@ -219,10 +219,10 @@ class LabelFilterBackend(filters.BaseFilterBackend):
         if not hasattr(queryset, 'model'):
             return queryset
 
-        if not hasattr(queryset.model, 'labels'):
+        if not hasattr(queryset.model, 'label_model'):
             return queryset
 
-        model = queryset.model
+        model = queryset.model.label_model()
         labeled_resource_cls = model._labels.field.related_model
         app_label = model._meta.app_label
         model_name = model._meta.model_name
diff --git a/apps/common/signal_handlers.py b/apps/common/signal_handlers.py
index ba5a4ad92..ad765657e 100644
--- a/apps/common/signal_handlers.py
+++ b/apps/common/signal_handlers.py
@@ -69,7 +69,7 @@ def digest_sql_query():
 
         for query in queries:
             sql = query['sql']
-            print("  # {}: {}".format(query['time'], sql))
+            print("  # {}: {}".format(query['time'], sql[:1000]))
         if len(queries) < 3:
             continue
         print("- Table: {}".format(table_name))
diff --git a/apps/labels/api.py b/apps/labels/api.py
index 44a6702e7..1c6965a43 100644
--- a/apps/labels/api.py
+++ b/apps/labels/api.py
@@ -73,7 +73,7 @@ class LabelContentTypeResourceViewSet(JMSModelViewSet):
             queryset = model.objects.all()
         if bound == '1':
             queryset = queryset.filter(id__in=list(res_ids))
-        elif bound == '0':
+        else:
             queryset = queryset.exclude(id__in=list(res_ids))
         keyword = self.request.query_params.get('search')
         if keyword:
@@ -90,9 +90,10 @@ class LabelContentTypeResourceViewSet(JMSModelViewSet):
         LabeledResource.objects \
             .filter(res_type=content_type, label=label) \
             .exclude(res_id__in=res_ids).delete()
-        resources = []
-        for res_id in res_ids:
-            resources.append(LabeledResource(res_type=content_type, res_id=res_id, label=label, org_id=current_org.id))
+        resources = [
+            LabeledResource(res_type=content_type, res_id=res_id, label=label, org_id=current_org.id)
+            for res_id in res_ids
+        ]
         LabeledResource.objects.bulk_create(resources, ignore_conflicts=True)
         return Response({"total": len(res_ids)})
 
@@ -129,15 +130,22 @@ class LabeledResourceViewSet(OrgBulkModelViewSet):
     }
     ordering_fields = ('res_type', 'date_created')
 
-    # Todo: 这里需要优化,查询 sql 太多
     def filter_search(self, queryset):
         keyword = self.request.query_params.get('search')
         if not keyword:
             return queryset
+        keyword = keyword.strip().lower()
         matched = []
-        for instance in queryset:
-            if keyword.lower() in str(instance.resource).lower():
-                matched.append(instance.id)
+        offset = 0
+        limit = 10000
+        while True:
+            page = queryset[offset:offset + limit]
+            if not page:
+                break
+            offset += limit
+            for instance in page:
+                if keyword in str(instance.resource).lower():
+                    matched.append(instance.id)
         return queryset.filter(id__in=matched)
 
     def get_queryset(self):
diff --git a/apps/labels/mixins.py b/apps/labels/mixins.py
index 90b76f3ed..9feb6f0f0 100644
--- a/apps/labels/mixins.py
+++ b/apps/labels/mixins.py
@@ -1,5 +1,6 @@
 from django.contrib.contenttypes.fields import GenericRelation
 from django.db import models
+from django.db.models import OneToOneField
 
 from .models import LabeledResource
 
@@ -12,10 +13,25 @@ class LabeledMixin(models.Model):
     class Meta:
         abstract = True
 
+    @classmethod
+    def label_model(cls):
+        pk_field = cls._meta.pk
+        model = cls
+        if isinstance(pk_field, OneToOneField):
+            model = pk_field.related_model
+        return model
+
+    @property
+    def real(self):
+        pk_field = self._meta.pk
+        if isinstance(pk_field, OneToOneField):
+            return getattr(self, pk_field.name)
+        return self
+
     @property
     def labels(self):
-        return self._labels
+        return self.real._labels
 
     @labels.setter
     def labels(self, value):
-        self._labels.set(value, bulk=False)
+        self.real._labels.set(value, bulk=False)