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

perf: 修改 platform 内置创建
pull/9666/head
老广 2023-02-21 15:50:58 +08:00 committed by GitHub
commit b5d814e9be
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 9 deletions

View File

@ -253,17 +253,20 @@ class AllTypes(ChoicesMixin):
return data return data
@classmethod @classmethod
def create_or_update_by_platform_data(cls, name, platform_data): def create_or_update_by_platform_data(cls, name, platform_data, platform_cls=None):
from assets.models import Platform, PlatformAutomation, PlatformProtocol # 不直接用 Platform 是因为可能在 migrations 中使用
from assets.models import Platform
if platform_cls is None:
platform_cls = Platform
automation_data = platform_data.pop('automation', {}) automation_data = platform_data.pop('automation', {})
protocols_data = platform_data.pop('protocols', []) 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 defaults=platform_data, name=name
) )
if not platform.automation: if not platform.automation:
automation = PlatformAutomation.objects.create() automation = platform_cls.automation.field.related_model.objects.create()
platform.automation = automation platform.automation = automation
platform.save() platform.save()
else: else:
@ -275,10 +278,13 @@ class AllTypes(ChoicesMixin):
platform.protocols.all().delete() platform.protocols.all().delete()
for p in protocols_data: for p in protocols_data:
p.pop('primary', None) p.pop('primary', None)
PlatformProtocol.objects.create(**p, platform=platform) platform.protocols.create(**p)
@classmethod @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") print("\n\tCreate internal platforms")
for category, type_cls in cls.category_types(): for category, type_cls in cls.category_types():
print("\t## Category: {}".format(category.label)) print("\t## Category: {}".format(category.label))
@ -311,7 +317,7 @@ class AllTypes(ChoicesMixin):
'automation': {**default_automation, **_automation}, 'automation': {**default_automation, **_automation},
'protocols': protocols_data '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 @classmethod
def update_user_create_platforms(cls, platform_cls): def update_user_create_platforms(cls, platform_cls):
@ -328,4 +334,4 @@ class AllTypes(ChoicesMixin):
for platform in user_platforms: for platform in user_platforms:
print("\t- Update platform: {}".format(platform.name)) print("\t- Update platform: {}".format(platform.name))
platform_data = cls.get_type_default_platform(platform.category, platform.type) 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)

View File

@ -5,7 +5,8 @@ from assets.const import AllTypes
def create_internal_platforms(apps, *args): 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): def update_user_platforms(apps, *args):