From cb82b53791d082f18e287fd14dcba702d042696f Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Fri, 11 Nov 2022 16:13:16 +0800
Subject: [PATCH] perf: automation celery task (#9046)

Co-authored-by: feng <1304903146@qq.com>
---
 apps/assets/api/automations/base.py    | 4 +---
 apps/assets/models/automations/base.py | 2 +-
 apps/assets/tasks/automation.py        | 4 +++-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/apps/assets/api/automations/base.py b/apps/assets/api/automations/base.py
index 845b721c8..1b480cbdc 100644
--- a/apps/assets/api/automations/base.py
+++ b/apps/assets/api/automations/base.py
@@ -5,7 +5,6 @@ from rest_framework import status, mixins, viewsets
 
 from orgs.mixins import generics
 from assets import serializers
-from assets.const import AutomationTypes
 from assets.tasks import execute_automation
 from assets.models import BaseAutomation, AutomationExecution
 from common.const.choices import Trigger
@@ -111,8 +110,7 @@ class AutomationExecutionViewSet(
         serializer.is_valid(raise_exception=True)
         automation = serializer.validated_data.get('automation')
         tp = serializer.validated_data.get('type')
-        model = AutomationTypes.get_type_model(tp)
         task = execute_automation.delay(
-            pid=automation.pk, trigger=Trigger.manual, model=model
+            pid=automation.pk, trigger=Trigger.manual, tp=tp
         )
         return Response({'task': task.id}, status=status.HTTP_201_CREATED)
diff --git a/apps/assets/models/automations/base.py b/apps/assets/models/automations/base.py
index fd9bc422b..9977f6830 100644
--- a/apps/assets/models/automations/base.py
+++ b/apps/assets/models/automations/base.py
@@ -47,7 +47,7 @@ class BaseAutomation(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin):
     def get_register_task(self):
         name = f"automation_{self.type}_strategy_period_{str(self.id)[:8]}"
         task = execute_automation.name
-        args = (str(self.id), Trigger.timing, self._meta.model)
+        args = (str(self.id), Trigger.timing, self.type)
         kwargs = {}
         return name, task, args, kwargs
 
diff --git a/apps/assets/tasks/automation.py b/apps/assets/tasks/automation.py
index 873e606b6..d750fa7cb 100644
--- a/apps/assets/tasks/automation.py
+++ b/apps/assets/tasks/automation.py
@@ -2,12 +2,14 @@ from celery import shared_task
 
 from orgs.utils import tmp_to_root_org, tmp_to_org
 from common.utils import get_logger, get_object_or_none
+from assets.const import AutomationTypes
 
 logger = get_logger(__file__)
 
 
 @shared_task(queue='ansible')
-def execute_automation(pid, trigger, model):
+def execute_automation(pid, trigger, tp):
+    model = AutomationTypes.get_model(tp)
     with tmp_to_root_org():
         instance = get_object_or_none(model, pk=pid)
     if not instance: