perf: 优化账号任务

apps/accounts/api/account/__init__.py

@@ -1,4 +1,3 @@
 from .account import *
-from .backup import *
+from .task import *
 from .template import *
-from .gathered_account import *

apps/accounts/api/account/account.py

@@ -1,12 +1,11 @@
 from django.shortcuts import get_object_or_404
 from rest_framework.decorators import action
-from rest_framework.generics import CreateAPIView, ListAPIView
+from rest_framework.generics import ListAPIView
 from rest_framework.response import Response
 
 from accounts import serializers
 from accounts.filters import AccountFilterSet
 from accounts.models import Account
-from accounts.tasks import verify_accounts_connectivity_task, push_accounts_to_assets_task
 from assets.models import Asset
 from authentication.const import ConfirmType
 from common.permissions import UserConfirmation
@@ -15,7 +14,7 @@ from orgs.mixins.api import OrgBulkModelViewSet
 
 __all__ = [
     'AccountViewSet', 'AccountSecretsViewSet',
-    'AccountsTaskCreateAPI', 'AccountHistoriesSecretAPI'
+    'AccountHistoriesSecretAPI'
 ]
 
 from rbac.permissions import RBACPermission
@@ -37,6 +36,7 @@ class AccountViewSet(OrgBulkModelViewSet):
     def su_from_accounts(self, request, *args, **kwargs):
         account_id = request.query_params.get('account')
         asset_id = request.query_params.get('asset')
+
         if account_id:
             account = get_object_or_404(Account, pk=account_id)
             accounts = account.get_su_from_accounts()
@@ -75,39 +75,3 @@ class AccountHistoriesSecretAPI(RecordViewLogMixin, ListAPIView):
 
     def get_queryset(self):
         return self.model.objects.filter(id=self.kwargs.get('pk'))
-
-
-class AccountsTaskCreateAPI(CreateAPIView):
-    serializer_class = serializers.AccountTaskSerializer
-    search_fields = AccountViewSet.search_fields
-    filterset_class = AccountViewSet.filterset_class
-
-    def check_permissions(self, request):
-        return request.user.has_perm('assets.test_assetconnectivity')
-
-    def get_accounts(self):
-        queryset = Account.objects.all()
-        queryset = self.filter_queryset(queryset)
-        return queryset
-
-    def perform_create(self, serializer):
-        data = serializer.validated_data
-        accounts = data.get('accounts')
-        account_ids = accounts.values_list('id', flat=True)
-        asset_ids = [account.asset_id for account in accounts]
-
-        if data['action'] == 'push':
-            task = push_accounts_to_assets_task.delay(account_ids, asset_ids)
-        else:
-            task = verify_accounts_connectivity_task.delay(account_ids, asset_ids)
-
-        data = getattr(serializer, '_data', {})
-        data["task"] = task.id
-        setattr(serializer, '_data', data)
-        return task
-
-    def get_exception_handler(self):
-        def handler(e, context):
-            return Response({"error": str(e)}, status=400)
-
-        return handler

apps/accounts/api/account/gathered_account.py

@@ -1,42 +0,0 @@
-from rest_framework import status
-from rest_framework.decorators import action
-from rest_framework.response import Response
-from django.utils.translation import ugettext_lazy as _
-
-from accounts import serializers
-from accounts.const import Source
-from accounts.models import GatheredAccount
-from accounts.filters import GatheredAccountFilterSet
-from orgs.mixins.api import OrgBulkModelViewSet
-
-__all__ = [
-    'GatheredAccountViewSet',
-]
-
-
-class GatheredAccountViewSet(OrgBulkModelViewSet):
-    model = GatheredAccount
-    search_fields = ('username',)
-    filterset_class = GatheredAccountFilterSet
-    serializer_classes = {
-        'default': serializers.GatheredAccountSerializer,
-    }
-    rbac_perms = {
-        'sync_account': 'assets.add_gatheredaccount',
-    }
-
-    @action(methods=['post'], detail=True, url_path='sync')
-    def sync_account(self, request, *args, **kwargs):
-        gathered_account = super().get_object()
-        asset = gathered_account.asset
-        username = gathered_account.username
-        accounts = asset.accounts.filter(username=username)
-        if accounts.exists():
-            accounts.update(source=Source.COLLECTED)
-        else:
-            asset.accounts.model.objects.create(
-                asset=asset, username=username,
-                name=f'{username}-{_("Collected")}',
-                source=Source.COLLECTED
-            )
-        return Response(status=status.HTTP_201_CREATED)

apps/accounts/api/account/task.py

@@ -0,0 +1,37 @@
+from rest_framework.generics import CreateAPIView
+from rest_framework.response import Response
+
+from accounts import serializers
+from accounts.tasks import verify_accounts_connectivity_task, push_accounts_to_assets_task
+
+__all__ = [
+    'AccountsTaskCreateAPI',
+]
+
+
+class AccountsTaskCreateAPI(CreateAPIView):
+    serializer_class = serializers.AccountTaskSerializer
+
+    def check_permissions(self, request):
+        return request.user.has_perm('assets.test_assetconnectivity')
+
+    def perform_create(self, serializer):
+        data = serializer.validated_data
+        accounts = data.get('accounts', [])
+        account_ids = [a.id for a in accounts]
+
+        if data['action'] == 'push':
+            task = push_accounts_to_assets_task.delay(account_ids)
+        else:
+            task = verify_accounts_connectivity_task.delay(account_ids)
+
+        data = getattr(serializer, '_data', {})
+        data["task"] = task.id
+        setattr(serializer, '_data', data)
+        return task
+
+    def get_exception_handler(self):
+        def handler(e, context):
+            return Response({"error": str(e)}, status=400)
+
+        return handler

apps/accounts/api/automations/__init__.py

@@ -1,3 +1,4 @@
+from .backup import *
 from .base import *
 from .change_secret import *
 from .gather_accounts import *

apps/accounts/api/automations/gather_accounts.py

@@ -1,13 +1,22 @@
 # -*- coding: utf-8 -*-
 #
+from django.utils.translation import ugettext_lazy as _
+from rest_framework import status
+from rest_framework.decorators import action
+from rest_framework.response import Response
+
 from accounts import serializers
 from accounts.const import AutomationTypes
+from accounts.const import Source
+from accounts.filters import GatheredAccountFilterSet
 from accounts.models import GatherAccountsAutomation
+from accounts.models import GatheredAccount
 from orgs.mixins.api import OrgBulkModelViewSet
 from .base import AutomationExecutionViewSet
 
 __all__ = [
-    'GatherAccountsAutomationViewSet', 'GatherAccountsExecutionViewSet'
+    'GatherAccountsAutomationViewSet', 'GatherAccountsExecutionViewSet',
+    'GatheredAccountViewSet'
 ]
 
 
@@ -31,3 +40,32 @@ class GatherAccountsExecutionViewSet(AutomationExecutionViewSet):
         queryset = super().get_queryset()
         queryset = queryset.filter(automation__type=self.tp)
         return queryset
+
+
+class GatheredAccountViewSet(OrgBulkModelViewSet):
+    model = GatheredAccount
+    search_fields = ('username',)
+    filterset_class = GatheredAccountFilterSet
+    serializer_classes = {
+        'default': serializers.GatheredAccountSerializer,
+    }
+    rbac_perms = {
+        'sync_account': 'assets.add_gatheredaccount',
+    }
+
+    @action(methods=['post'], detail=True, url_path='sync')
+    def sync_account(self, request, *args, **kwargs):
+        gathered_account = super().get_object()
+        asset = gathered_account.asset
+        username = gathered_account.username
+        accounts = asset.accounts.filter(username=username)
+
+        if accounts.exists():
+            accounts.update(source=Source.COLLECTED)
+        else:
+            asset.accounts.model.objects.create(
+                asset=asset, username=username,
+                name=f'{username}-{_("Collected")}',
+                source=Source.COLLECTED
+            )
+        return Response(status=status.HTTP_201_CREATED)