jumpserver/apps/ops/mixin.py

# -*- coding: utf-8 -*-
#
import abc
from datetime import timedelta

from celery.schedules import crontab
from django.db import models
from django.utils import timezone
from django.utils.translation import gettext_lazy as _
from django_celery_beat.models import CrontabSchedule, IntervalSchedule, ClockedSchedule
from rest_framework import serializers

from .celery.utils import (
    create_or_update_celery_periodic_tasks, disable_celery_periodic_task,
    delete_celery_periodic_task,
)

__all__ = [
    'PeriodTaskModelMixin', 'PeriodTaskSerializerMixin',
]


class PeriodTaskModelQuerySet(models.QuerySet):
    def delete(self, *args, **kwargs):
        for obj in self:
            obj.delete()
        return super().delete(*args, **kwargs)


class PeriodTaskModelMixin(models.Model):
    name = models.CharField(
        max_length=128, unique=False, verbose_name=_("Name")
    )
    is_periodic = models.BooleanField(default=False, verbose_name=_("Periodic run"))
    interval = models.IntegerField(
        default=24, null=True, blank=True, verbose_name=_("Interval"),
    )
    crontab = models.CharField(
        blank=True, max_length=128, default='', verbose_name=_("Crontab"),
    )
    start_time = models.DateTimeField(
        blank=True, null=True,
        verbose_name=_('Start Datetime'),
        help_text=_(
            'Datetime when the schedule should begin '
            'triggering the task to run'
        ),
    )
    date_last_run = models.DateTimeField(blank=True, null=True, verbose_name=_("Date last run"))
    objects = PeriodTaskModelQuerySet.as_manager()

    @abc.abstractmethod
    def get_register_task(self):
        period_name, task, args, kwargs = None, None, (), {}
        return period_name, task, args, kwargs

    @property
    def interval_ratio(self):
        return 3600, 'h'

    @property
    def interval_display(self):
        unit = self.interval_ratio[1]
        if not self.interval:
            return ''
        return '{} {}'.format(self.interval, unit)

    def set_period_schedule(self):
        name, task, args, kwargs = self.get_register_task()
        is_active = self.is_active if hasattr(self, 'is_active') else True
        if not self.is_periodic or not is_active:
            disable_celery_periodic_task(name)
            return

        cron = interval = None
        if self.crontab:
            cron = self.crontab
        elif self.interval:
            interval = self.interval * self.interval_ratio[0]

        tasks = {
            name: {
                'task': task,
                'interval': interval,
                'crontab': cron,
                'args': args,
                'kwargs': kwargs,
                'enabled': True,
                'start_time': self.start_time,
            }
        }
        create_or_update_celery_periodic_tasks(tasks)

    def execute(self, *args, **kwargs):
        self.date_last_run = timezone.now()

    def save(self, *args, **kwargs):
        instance = super().save(**kwargs)
        self.set_period_schedule()
        return instance

    def delete(self, using=None, keep_parents=False):
        name = self.get_register_task()[0]
        instance = super().delete(using=using, keep_parents=keep_parents)
        delete_celery_periodic_task(name)
        return instance

    @property
    def periodic_display(self):
        if self.is_periodic and self.crontab:
            return _('Crontab') + " ( {} )".format(self.crontab)
        if self.is_periodic and self.interval:
            return _('Interval') + " ( {} h )".format(self.interval)
        return '-'

    @property
    def schedule(self):
        from django_celery_beat.models import PeriodicTask
        name = self.get_register_task()[0]
        return PeriodicTask.objects.filter(name=name).first()

    def get_next_run_time(self):
        if not self.is_periodic:
            return None
        task = self.schedule
        now = task.schedule.nowfun()

        if self.start_time and self.start_time > now:
            return self.start_time

        scheduler = task.scheduler
        # 根据不同的调度类型计算下次执行时间
        if isinstance(scheduler, CrontabSchedule):
            schedule = crontab(
                minute=scheduler.minute,
                hour=scheduler.hour,
                day_of_week=scheduler.day_of_week,
                day_of_month=scheduler.day_of_month,
                month_of_year=scheduler.month_of_year,
            )
            next_run = schedule.remaining_estimate(now)
            return now + next_run
        elif isinstance(scheduler, IntervalSchedule):
            interval = timedelta(
                seconds=scheduler.every * {
                    IntervalSchedule.SECONDS: 1,
                    IntervalSchedule.MINUTES: 60,
                    IntervalSchedule.HOURS: 3600,
                    IntervalSchedule.DAYS: 86400,
                }[scheduler.period]
            )
            last_run = task.last_run_at or now
            return last_run + interval

        elif isinstance(scheduler, ClockedSchedule):
            return scheduler.clocked_time

        else:
            raise ValueError("不支持的任务调度类型")

    class Meta:
        abstract = True


class PeriodTaskSerializerMixin(serializers.Serializer):
    is_periodic = serializers.BooleanField(default=True, label=_("Periodic run"))
    crontab = serializers.CharField(
        max_length=128, allow_blank=True,
        allow_null=True, required=False, label=_('Crontab')
    )
    interval = serializers.IntegerField(
        default=24, allow_null=True, required=False, label=_('Interval'),
        max_value=65535, min_value=1,
    )
    start_time = serializers.DateTimeField(
        allow_null=True, required=False,
        label=_('Start Datetime'),
        help_text=_(
            'Datetime when the schedule should begin '
            'triggering the task to run'
        ),
    )
    periodic_display = serializers.CharField(read_only=True, label=_('Run period'))

    def validate_crontab(self, crontab):
        if not crontab:
            return crontab
        if isinstance(crontab, str) and len(crontab.strip().split()) != 5:
            msg = _('* Please enter a valid crontab expression')
            raise serializers.ValidationError(msg)

        crontab = crontab.strip().split()
        if '*' in crontab[0]:
            raise serializers.ValidationError(_("Crontab minute must not contain '*'"))
        return crontab

    def validate_interval(self, interval):
        if not interval and not isinstance(interval, int):
            return interval
        return interval

    def validate_is_periodic(self, ok):
        if not ok:
            return ok
        crontab = self.initial_data.get('crontab')
        interval = self.initial_data.get('interval')
        if ok and not any([crontab, interval]):
            msg = _("Require interval or crontab setting")
            raise serializers.ValidationError(msg)
        return ok

    def validate(self, attrs):
        attrs = super().validate(attrs)
        if not attrs.get('is_periodic'):
            attrs['interval'] = None
            attrs['crontab'] = ''
        if attrs.get('crontab'):
            attrs['interval'] = None
        return attrs