Merge pull request #9623 from jumpserver/pr@dev@account_push

perf: account push

apps/accounts/api/automations/base.py

@@ -110,6 +110,6 @@ class AutomationExecutionViewSet(
         serializer.is_valid(raise_exception=True)
         automation = serializer.validated_data.get('automation')
         task = execute_automation.delay(
-            automation=automation, trigger=Trigger.manual
+            pid=str(automation.pk), trigger=Trigger.manual, tp=self.tp
         )
         return Response({'task': task.id}, status=status.HTTP_201_CREATED)

apps/accounts/tasks/automation.py

@@ -1,21 +1,35 @@
 from celery import shared_task
 from django.utils.translation import gettext_lazy as _
 
-from common.utils import get_logger
-from orgs.utils import tmp_to_org
+from accounts.const import AutomationTypes
+from common.utils import get_logger, get_object_or_none
+from orgs.utils import tmp_to_org, tmp_to_root_org
 
 logger = get_logger(__file__)
 
 
-def task_activity_callback(self, instance, *args):
-    asset_ids = instance.get_all_asset_ids()
-    return asset_ids, instance.org_id
+def task_activity_callback(self, pid, trigger, tp):
+    model = AutomationTypes.get_type_model(tp)
+    with tmp_to_root_org():
+        instance = get_object_or_none(model, pk=pid)
+    if not instance:
+        return
+    if not instance.latest_execution:
+        return
+    resource_ids = instance.latest_execution.get_all_asset_ids()
+    return resource_ids, instance.org_id
 
 
 @shared_task(
     queue='ansible', verbose_name=_('Account execute automation'),
     activity_callback=task_activity_callback
 )
-def execute_automation(instance, trigger):
+def execute_automation(pid, trigger, tp):
+    model = AutomationTypes.get_type_model(tp)
+    with tmp_to_root_org():
+        instance = get_object_or_none(model, pk=pid)
+    if not instance:
+        logger.error("No automation task found: {}".format(pid))
+        return
     with tmp_to_org(instance.org):
         instance.execute(trigger)