Merge pull request #9665 from jumpserver/pr@dev@fix_platform_internal

perf: 修改 platform 内置创建

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)

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):