From 771ea174695ecb8efdbcd1e15940a7e0c104a5f1 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Tue, 21 Feb 2023 15:49:59 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20platform=20?=
 =?UTF-8?q?=E5=86=85=E7=BD=AE=E5=88=9B=E5=BB=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/assets/const/types.py                    | 22 ++++++++++++-------
 .../migrations/0097_auto_20220426_1558.py     |  3 ++-
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/apps/assets/const/types.py b/apps/assets/const/types.py
index 509c2f58e..3c0989517 100644
--- a/apps/assets/const/types.py
+++ b/apps/assets/const/types.py
@@ -253,17 +253,20 @@ class AllTypes(ChoicesMixin):
         return data
 
     @classmethod
-    def create_or_update_by_platform_data(cls, name, platform_data):
-        from assets.models import Platform, PlatformAutomation, PlatformProtocol
+    def create_or_update_by_platform_data(cls, name, platform_data, platform_cls=None):
+        # 不直接用 Platform 是因为可能在 migrations 中使用
+        from assets.models import Platform
+        if platform_cls is None:
+            platform_cls = Platform
 
         automation_data = platform_data.pop('automation', {})
         protocols_data = platform_data.pop('protocols', [])
 
-        platform, created = Platform.objects.update_or_create(
+        platform, created = platform_cls.objects.update_or_create(
             defaults=platform_data, name=name
         )
         if not platform.automation:
-            automation = PlatformAutomation.objects.create()
+            automation = platform_cls.automation.field.related_model.objects.create()
             platform.automation = automation
             platform.save()
         else:
@@ -275,10 +278,13 @@ class AllTypes(ChoicesMixin):
         platform.protocols.all().delete()
         for p in protocols_data:
             p.pop('primary', None)
-            PlatformProtocol.objects.create(**p, platform=platform)
+            platform.protocols.create(**p)
 
     @classmethod
-    def create_or_update_internal_platforms(cls):
+    def create_or_update_internal_platforms(cls, platform_cls=None):
+        if platform_cls is None:
+            platform_cls = cls
+
         print("\n\tCreate internal platforms")
         for category, type_cls in cls.category_types():
             print("\t## Category: {}".format(category.label))
@@ -311,7 +317,7 @@ class AllTypes(ChoicesMixin):
                         'automation': {**default_automation, **_automation},
                         'protocols': protocols_data
                     }
-                    cls.create_or_update_by_platform_data(name, platform_data)
+                    cls.create_or_update_by_platform_data(name, platform_data, platform_cls=platform_cls)
 
     @classmethod
     def update_user_create_platforms(cls, platform_cls):
@@ -328,4 +334,4 @@ class AllTypes(ChoicesMixin):
         for platform in user_platforms:
             print("\t- Update platform: {}".format(platform.name))
             platform_data = cls.get_type_default_platform(platform.category, platform.type)
-            cls.create_or_update_by_platform_data(platform.name, platform_data)
+            cls.create_or_update_by_platform_data(platform.name, platform_data, platform_cls=platform_cls)
diff --git a/apps/assets/migrations/0097_auto_20220426_1558.py b/apps/assets/migrations/0097_auto_20220426_1558.py
index 2a16841d5..59d9726e7 100644
--- a/apps/assets/migrations/0097_auto_20220426_1558.py
+++ b/apps/assets/migrations/0097_auto_20220426_1558.py
@@ -5,7 +5,8 @@ from assets.const import AllTypes
 
 
 def create_internal_platforms(apps, *args):
-    AllTypes.create_or_update_internal_platforms()
+    platform_cls = apps.get_model('assets', 'Platform')
+    AllTypes.create_or_update_internal_platforms(platform_cls)
 
 
 def update_user_platforms(apps, *args):