feat: 丰富资产任务API创建；支持针对多个系统用户一个资产的推送和测试

apps/assets/api/asset.py

@@ -9,10 +9,11 @@ from common.utils import get_logger, get_object_or_none
 from common.permissions import IsOrgAdmin, IsOrgAdminOrAppUser, IsSuperUser
 from orgs.mixins.api import OrgBulkModelViewSet
 from orgs.mixins import generics
-from ..models import Asset, Node, Platform
+from ..models import Asset, Node, Platform, SystemUser
 from .. import serializers
 from ..tasks import (
-    update_assets_hardware_info_manual, test_assets_connectivity_manual
+    update_assets_hardware_info_manual, test_assets_connectivity_manual,
+    test_system_users_connectivity_a_asset, push_system_users_a_asset
 )
 from ..filters import FilterAssetByNodeFilterBackend, LabelFilterBackend, IpInFilterBackend
 
@@ -94,21 +95,27 @@ class AssetPlatformViewSet(ModelViewSet):
 
 
 class AssetsTaskMixin:
+
     def perform_assets_task(self, serializer):
         data = serializer.validated_data
-        assets = data['assets']
         action = data['action']
+        assets = data.get('assets', [])
         if action == "refresh":
             task = update_assets_hardware_info_manual.delay(assets)
         else:
+            # action == 'test':
             task = test_assets_connectivity_manual.delay(assets)
+        return task
+
+    def perform_create(self, serializer):
+        task = self.perform_assets_task(serializer)
+        self.set_task_to_serializer_data(serializer, task)
+
+    def set_task_to_serializer_data(self, serializer, task):
         data = getattr(serializer, '_data', {})
         data["task"] = task.id
         setattr(serializer, '_data', data)
 
-    def perform_create(self, serializer):
-        self.perform_assets_task(serializer)
-
 
 class AssetTaskCreateApi(AssetsTaskMixin, generics.CreateAPIView):
     model = Asset
@@ -117,13 +124,37 @@ class AssetTaskCreateApi(AssetsTaskMixin, generics.CreateAPIView):
 
     def create(self, request, *args, **kwargs):
         pk = self.kwargs.get('pk')
+        request.data['asset'] = pk
         request.data['assets'] = [pk]
         return super().create(request, *args, **kwargs)
 
+    def perform_asset_task(self, serializer):
+        data = serializer.validated_data
+        action = data['action']
+        if action not in ['push_system_user', 'test_system_user']:
+            return
+        asset = data['asset']
+        system_users = data.get('system_users')
+        if not system_users:
+            system_users = asset.get_all_systemusers()
+        if action == 'push_system_user':
+            task = push_system_users_a_asset.delay(system_users, asset=asset)
+        elif action == 'test_system_user':
+            task = test_system_users_connectivity_a_asset.delay(system_users, asset=asset)
+        else:
+            task = None
+        return task
+
+    def perform_create(self, serializer):
+        task = self.perform_asset_task(serializer)
+        if not task:
+            task = self.perform_assets_task(serializer)
+        self.set_task_to_serializer_data(serializer, task)
+
 
 class AssetsTaskCreateApi(AssetsTaskMixin, generics.CreateAPIView):
     model = Asset
-    serializer_class = serializers.AssetTaskSerializer
+    serializer_class = serializers.AssetsTaskSerializer
     permission_classes = (IsOrgAdmin,)
 
 

apps/assets/models/asset.py

@@ -345,6 +345,13 @@ class Asset(AbsConnectivity, ProtocolsMixin, NodesRelationMixin, OrgModelMixin):
         tree_node = TreeNode(**data)
         return tree_node
 
+    def get_all_systemusers(self):
+        from .user import SystemUser
+        system_user_ids = SystemUser.assets.through.objects.filter(asset=self)\
+            .values_list('systemuser_id', flat=True)
+        system_users = SystemUser.objects.filter(id__in=system_user_ids)
+        return system_users
+
     class Meta:
         unique_together = [('org_id', 'hostname')]
         verbose_name = _("Asset")

apps/assets/serializers/asset.py

@@ -10,7 +10,7 @@ from ..models import Asset, Node, Platform, SystemUser
 __all__ = [
     'AssetSerializer', 'AssetSimpleSerializer',
     'ProtocolsField', 'PlatformSerializer',
-    'AssetTaskSerializer',
+    'AssetTaskSerializer', 'AssetsTaskSerializer'
 ]
 
 
@@ -183,7 +183,7 @@ class AssetSimpleSerializer(serializers.ModelSerializer):
         fields = ['id', 'hostname', 'ip', 'port', 'connectivity', 'date_verified']
 
 
-class AssetTaskSerializer(serializers.Serializer):
+class AssetsTaskSerializer(serializers.Serializer):
     ACTION_CHOICES = (
         ('refresh', 'refresh'),
         ('test', 'test'),
@@ -193,3 +193,16 @@ class AssetTaskSerializer(serializers.Serializer):
     assets = serializers.PrimaryKeyRelatedField(
         queryset=Asset.objects, required=False, allow_empty=True, many=True
     )
+
+
+class AssetTaskSerializer(AssetsTaskSerializer):
+    ACTION_CHOICES = tuple(list(AssetsTaskSerializer.ACTION_CHOICES) + [
+        ('push_system_user', 'push_system_user'),
+        ('test_system_user', 'test_system_user')
+    ])
+    asset = serializers.PrimaryKeyRelatedField(
+        queryset=Asset.objects, required=False, allow_empty=True, many=False
+    )
+    system_users = serializers.PrimaryKeyRelatedField(
+        queryset=SystemUser.objects, required=False, allow_empty=True, many=True
+    )

apps/assets/tasks/push_system_user.py

@@ -17,6 +17,7 @@ logger = get_logger(__file__)
 __all__ = [
     'push_system_user_util', 'push_system_user_to_assets',
     'push_system_user_to_assets_manual', 'push_system_user_a_asset_manual',
+    'push_system_users_a_asset'
 ]
 
 
@@ -280,14 +281,21 @@ def push_system_user_a_asset_manual(system_user, asset, username=None):
     """
     将系统用户推送到一个资产上
     """
-    if username is None:
+    # if username is None:
-        username = system_user.username
+    #     username = system_user.username
     task_name = _("Push system users to asset: {}({}) => {}").format(
         system_user.name, username, asset
     )
     return push_system_user_util(system_user, [asset], task_name=task_name, username=username)
 
 
+@shared_task(queue="ansible")
+@tmp_to_root_org()
+def push_system_users_a_asset(system_users, asset):
+    for system_user in system_users:
+        push_system_user_a_asset_manual(system_user, asset)
+
+
 @shared_task(queue="ansible")
 @tmp_to_root_org()
 def push_system_user_to_assets(system_user_id, asset_ids, username=None):

apps/assets/tasks/system_user_connectivity.py

@@ -18,6 +18,7 @@ logger = get_logger(__name__)
 __all__ = [
     'test_system_user_connectivity_util', 'test_system_user_connectivity_manual',
     'test_system_user_connectivity_period', 'test_system_user_connectivity_a_asset',
+    'test_system_users_connectivity_a_asset'
 ]
 
 
@@ -131,6 +132,12 @@ def test_system_user_connectivity_a_asset(system_user, asset):
     test_system_user_connectivity_util(system_user, [asset], task_name)
 
 
+@shared_task(queue="ansible")
+def test_system_users_connectivity_a_asset(system_users, asset):
+    for system_user in system_users:
+        test_system_user_connectivity_a_asset(system_user, asset)
+
+
 @shared_task(queue="ansible")
 def test_system_user_connectivity_period():
     if not const.PERIOD_TASK_ENABLED: