diff --git a/apps/accounts/api/automations/base.py b/apps/accounts/api/automations/base.py index f4aec2093..700484e9b 100644 --- a/apps/accounts/api/automations/base.py +++ b/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) diff --git a/apps/accounts/tasks/automation.py b/apps/accounts/tasks/automation.py index 27758c0be..bed1de77c 100644 --- a/apps/accounts/tasks/automation.py +++ b/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)