From 93da3e58f26c51ddb52ea49ae1eff93a08e93f61 Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Tue, 12 Dec 2023 14:18:26 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E3=80=90=E4=BC=98=E5=8C=96=E7=B3=BB?=
 =?UTF-8?q?=E7=BB=9F=E4=BB=BB=E5=8A=A1=E3=80=91=E6=94=AF=E6=8C=81=E6=98=BE?=
 =?UTF-8?q?=E7=A4=BA=20=E6=89=A7=E8=A1=8C=E5=91=A8=E6=9C=9F=E3=80=81?=
 =?UTF-8?q?=E4=B8=8B=E6=AC=A1=E5=BC=80=E5=A7=8B=E6=97=B6=E9=97=B4=20?=
 =?UTF-8?q?=E5=AD=97=E6=AE=B5=20(#12298)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: feng <1304903146@qq.com>
---
 apps/ops/api/celery.py         | 45 ++++++++++++++++++++++++++++++++--
 apps/ops/serializers/celery.py |  8 +++++-
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/apps/ops/api/celery.py b/apps/ops/api/celery.py
index a7baaf6e4..797d00c5e 100644
--- a/apps/ops/api/celery.py
+++ b/apps/ops/api/celery.py
@@ -7,14 +7,15 @@ from celery.result import AsyncResult
 from django.shortcuts import get_object_or_404
 from django.utils.translation import gettext as _
 from django_celery_beat.models import PeriodicTask
+from django_filters import rest_framework as drf_filters
 from rest_framework import generics, viewsets, mixins, status
 from rest_framework.response import Response
-from django_filters import rest_framework as drf_filters
 
 from common.api import LogTailApi, CommonApiMixin
+from common.drf.filters import BaseFilterSet
 from common.exceptions import JMSException
 from common.permissions import IsValidUser
-from common.drf.filters import BaseFilterSet
+from common.utils.timezone import local_now
 from ops.celery import app
 from ..ansible.utils import get_ansible_task_log_path
 from ..celery.utils import get_celery_task_log_path
@@ -138,6 +139,46 @@ class CeleryTaskViewSet(
     def get_queryset(self):
         return CeleryTask.objects.exclude(name__startswith='celery')
 
+    @staticmethod
+    def extract_schedule(input_string):
+        pattern = r'(\S+ \S+ \S+ \S+ \S+).*'
+        match = re.match(pattern, input_string)
+        if match:
+            return match.group(1)
+        else:
+            return input_string
+
+    def generate_execute_time(self, queryset):
+        names = [i.name for i in queryset]
+        periodic_tasks = PeriodicTask.objects.filter(name__in=names)
+        periodic_task_dict = {task.name: task for task in periodic_tasks}
+        now = local_now()
+        for i in queryset:
+            task = periodic_task_dict.get(i.name)
+            if not task:
+                continue
+            i.exec_cycle = self.extract_schedule(str(task.scheduler))
+
+            last_run_at = task.last_run_at or now
+            next_run_at = task.schedule.remaining_estimate(last_run_at)
+            if next_run_at.total_seconds() < 0:
+                next_run_at = task.schedule.remaining_estimate(now)
+            i.next_exec_time = now + next_run_at
+        return queryset
+
+    def list(self, request, *args, **kwargs):
+        queryset = self.filter_queryset(self.get_queryset())
+
+        page = self.paginate_queryset(queryset)
+        if page is not None:
+            page = self.generate_execute_time(page)
+            serializer = self.get_serializer(page, many=True)
+            return self.get_paginated_response(serializer.data)
+
+        queryset = self.generate_execute_time(queryset)
+        serializer = self.get_serializer(queryset, many=True)
+        return Response(serializer.data)
+
 
 class CeleryTaskExecutionViewSet(CommonApiMixin, viewsets.ModelViewSet):
     serializer_class = CeleryTaskExecutionSerializer
diff --git a/apps/ops/serializers/celery.py b/apps/ops/serializers/celery.py
index 5166aa755..a4e9c7e8d 100644
--- a/apps/ops/serializers/celery.py
+++ b/apps/ops/serializers/celery.py
@@ -30,9 +30,15 @@ class CeleryPeriodTaskSerializer(serializers.ModelSerializer):
 
 
 class CeleryTaskSerializer(serializers.ModelSerializer):
+    exec_cycle = serializers.CharField(read_only=True)
+    next_exec_time = serializers.DateTimeField(format="%Y/%m/%d %H:%M:%S", read_only=True)
+
     class Meta:
         model = CeleryTask
-        read_only_fields = ['id', 'name', 'meta', 'summary', 'state', 'date_last_publish']
+        read_only_fields = [
+            'id', 'name', 'meta', 'summary', 'state',
+            'date_last_publish', 'exec_cycle', 'next_exec_time'
+        ]
         fields = read_only_fields