[Update] 迁移celery到ops

apps/ops/__init__.py

@@ -1 +1 @@
-
+from .celery import app as celery_app

apps/ops/celery/__init__.py

@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+
+import os
+
+from celery import Celery
+
+# set the default Django settings module for the 'celery' program.
+os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'jumpserver.settings')
+
+from django.conf import settings
+
+app = Celery('jumpserver')
+
+# Using a string here means the worker will not have to
+# pickle the object when using Windows.
+app.config_from_object('django.conf:settings', namespace='CELERY')
+app.autodiscover_tasks(lambda: [app_config.split('.')[0] for app_config in settings.INSTALLED_APPS])
+

apps/ops/celery/const.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+#
+

apps/ops/celery/signal_handler.py

@@ -0,0 +1,89 @@
+# -*- coding: utf-8 -*-
+#
+import os
+import datetime
+import sys
+
+from django.conf import settings
+from django.core.cache import cache
+from celery import subtask
+from celery.signals import worker_ready, worker_shutdown, task_prerun, \
+    task_postrun, after_task_publish
+
+from django_celery_beat.models import PeriodicTask
+
+from common.utils import get_logger, TeeObj
+from common.const import celery_task_pre_key
+
+from .utils import get_after_app_ready_tasks, get_after_app_shutdown_clean_tasks
+
+
+logger = get_logger(__file__)
+
+WAITING = "waiting"
+RUNNING = "running"
+FINISHED = "finished"
+
+EXPIRE_TIME = 3600
+
+
+@worker_ready.connect
+def on_app_ready(sender=None, headers=None, body=None, **kwargs):
+    if cache.get("CELERY_APP_READY", 0) == 1:
+        return
+    cache.set("CELERY_APP_READY", 1, 10)
+    logger.debug("App ready signal recv")
+    tasks = get_after_app_ready_tasks()
+    logger.debug("Start need start task: [{}]".format(
+        ", ".join(tasks))
+    )
+    for task in tasks:
+        subtask(task).delay()
+
+
+@worker_shutdown.connect
+def after_app_shutdown(sender=None, headers=None, body=None, **kwargs):
+    if cache.get("CELERY_APP_SHUTDOWN", 0) == 1:
+        return
+    cache.set("CELERY_APP_SHUTDOWN", 1, 10)
+    tasks = get_after_app_shutdown_clean_tasks()
+    logger.debug("App shutdown signal recv")
+    logger.debug("Clean need cleaned period tasks: [{}]".format(
+        ', '.join(tasks))
+    )
+    PeriodicTask.objects.filter(name__in=tasks).delete()
+
+
+@task_prerun.connect
+def pre_run_task_signal_handler(sender, task_id=None, task=None, **kwargs):
+    task_key = celery_task_pre_key + task_id
+    info = cache.get(task_key, {})
+    now = datetime.datetime.now().strftime("%Y-%m-%d")
+    log_dir = os.path.join(settings.PROJECT_DIR, "data", "celery", now)
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    log_path = os.path.join(log_dir, task_id + '.log')
+    info.update({"status": RUNNING, "log_path": log_path})
+    cache.set(task_key, info, EXPIRE_TIME)
+    f = open(log_path, 'w')
+    tee = TeeObj(f)
+    sys.stdout = tee
+    task.log_f = tee
+
+
+@task_postrun.connect
+def post_run_task_signal_handler(sender, task_id=None, task=None, **kwargs):
+    task_key = celery_task_pre_key + task_id
+    info = cache.get(task_key, {})
+    info.update({"status": FINISHED})
+    cache.set(task_key, info, EXPIRE_TIME)
+    task.log_f.flush()
+    sys.stdout = task.log_f.origin_stdout
+    task.log_f.close()
+
+
+@after_task_publish.connect
+def after_task_publish_signal_handler(sender, headers=None, **kwargs):
+    task_id = headers["id"]
+    key = celery_task_pre_key + task_id
+    cache.set(key, {"status": WAITING}, EXPIRE_TIME)

apps/ops/celery/utils.py

@@ -0,0 +1,185 @@
+# -*- coding: utf-8 -*-
+#
+import json
+from functools import wraps
+
+from django.db.utils import ProgrammingError, OperationalError
+from django.core.cache import cache
+from django_celery_beat.models import PeriodicTask, IntervalSchedule, CrontabSchedule
+
+
+def add_register_period_task(name):
+    key = "__REGISTER_PERIODIC_TASKS"
+    value = cache.get(key, [])
+    value.append(name)
+    cache.set(key, value)
+
+
+def get_register_period_tasks():
+    key = "__REGISTER_PERIODIC_TASKS"
+    return cache.get(key, [])
+
+
+def add_after_app_shutdown_clean_task(name):
+    key = "__AFTER_APP_SHUTDOWN_CLEAN_TASKS"
+    value = cache.get(key, [])
+    value.append(name)
+    cache.set(key, value)
+
+
+def get_after_app_shutdown_clean_tasks():
+    key = "__AFTER_APP_SHUTDOWN_CLEAN_TASKS"
+    return cache.get(key, [])
+
+
+def add_after_app_ready_task(name):
+    key = "__AFTER_APP_READY_RUN_TASKS"
+    value = cache.get(key, [])
+    value.append(name)
+    cache.set(key, value)
+
+
+def get_after_app_ready_tasks():
+    key = "__AFTER_APP_READY_RUN_TASKS"
+    return cache.get(key, [])
+
+
+def create_or_update_celery_periodic_tasks(tasks):
+    """
+    :param tasks: {
+        'add-every-monday-morning': {
+            'task': 'tasks.add' # A registered celery task,
+            'interval': 30,
+            'crontab': "30 7 * * *",
+            'args': (16, 16),
+            'kwargs': {},
+            'enabled': False,
+        },
+    }
+    :return:
+    """
+    # Todo: check task valid, task and callback must be a celery task
+    for name, detail in tasks.items():
+        interval = None
+        crontab = None
+        try:
+            IntervalSchedule.objects.all().count()
+        except (ProgrammingError, OperationalError):
+            return None
+
+        if isinstance(detail.get("interval"), int):
+            intervals = IntervalSchedule.objects.filter(
+                every=detail["interval"], period=IntervalSchedule.SECONDS
+            )
+            if intervals:
+                interval = intervals[0]
+            else:
+                interval = IntervalSchedule.objects.create(
+                    every=detail['interval'],
+                    period=IntervalSchedule.SECONDS,
+                )
+        elif isinstance(detail.get("crontab"), str):
+            try:
+                minute, hour, day, month, week = detail["crontab"].split()
+            except ValueError:
+                raise SyntaxError("crontab is not valid")
+            kwargs = dict(
+                minute=minute, hour=hour, day_of_week=week,
+                day_of_month=day, month_of_year=month,
+            )
+            contabs = CrontabSchedule.objects.filter(
+                **kwargs
+            )
+            if contabs:
+                crontab = contabs[0]
+            else:
+                crontab = CrontabSchedule.objects.create(**kwargs)
+        else:
+            raise SyntaxError("Schedule is not valid")
+
+        defaults = dict(
+            interval=interval,
+            crontab=crontab,
+            name=name,
+            task=detail['task'],
+            args=json.dumps(detail.get('args', [])),
+            kwargs=json.dumps(detail.get('kwargs', {})),
+            enabled=detail.get('enabled', True),
+        )
+
+        task = PeriodicTask.objects.update_or_create(
+            defaults=defaults, name=name,
+        )
+        return task
+
+
+def disable_celery_periodic_task(task_name):
+    from django_celery_beat.models import PeriodicTask
+    PeriodicTask.objects.filter(name=task_name).update(enabled=False)
+
+
+def delete_celery_periodic_task(task_name):
+    from django_celery_beat.models import PeriodicTask
+    PeriodicTask.objects.filter(name=task_name).delete()
+
+
+def register_as_period_task(crontab=None, interval=None):
+    """
+    Warning: Task must be have not any args and kwargs
+    :param crontab:  "* * * * *"
+    :param interval:  60*60*60
+    :return:
+    """
+    if crontab is None and interval is None:
+        raise SyntaxError("Must set crontab or interval one")
+
+    def decorate(func):
+        if crontab is None and interval is None:
+            raise SyntaxError("Interval and crontab must set one")
+
+        # Because when this decorator run, the task was not created,
+        # So we can't use func.name
+        name = '{func.__module__}.{func.__name__}'.format(func=func)
+        if name not in get_register_period_tasks():
+            create_or_update_celery_periodic_tasks({
+                name: {
+                    'task': name,
+                    'interval': interval,
+                    'crontab': crontab,
+                    'args': (),
+                    'enabled': True,
+                }
+            })
+            add_register_period_task(name)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        return wrapper
+    return decorate
+
+
+def after_app_ready_start(func):
+    # Because when this decorator run, the task was not created,
+    # So we can't use func.name
+    name = '{func.__module__}.{func.__name__}'.format(func=func)
+    if name not in get_after_app_ready_tasks():
+        add_after_app_ready_task(name)
+
+    @wraps(func)
+    def decorate(*args, **kwargs):
+        return func(*args, **kwargs)
+    return decorate
+
+
+def after_app_shutdown_clean(func):
+    # Because when this decorator run, the task was not created,
+    # So we can't use func.name
+    name = '{func.__module__}.{func.__name__}'.format(func=func)
+    if name not in get_after_app_shutdown_clean_tasks():
+        add_after_app_shutdown_clean_task(name)
+
+    @wraps(func)
+    def decorate(*args, **kwargs):
+        return func(*args, **kwargs)
+    return decorate

apps/ops/models.py

@@ -15,7 +15,7 @@ from django_celery_beat.models import CrontabSchedule, IntervalSchedule, \
     PeriodicTask
 
 from common.utils import get_signer, get_logger
-from common.celery import delete_celery_periodic_task, \
+from .celery.utils import delete_celery_periodic_task, \
     create_or_update_celery_periodic_tasks, \
     disable_celery_periodic_task
 from .ansible import AdHocRunner, AnsibleError
@@ -218,14 +218,12 @@ class AdHoc(models.Model):
             hid = str(uuid.uuid4())
         history = AdHocRunHistory(id=hid, adhoc=self, task=self.task)
         time_start = time.time()
-        # f = open(history.log_path, 'w')
         try:
-            date_start = timezone.now().strftime('%Y-%m-%d %H:%M:%S')
-            # f.write("{} {}\r\n\r\n".format(date_start, self.task.name))
+            date_start = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+            print("{} Start task: {}\r\n".format(date_start, self.task.name))
             raw, summary = self._run_only()
-            # raw, summary = self._run_only(file_obj=f)
-            date_end = timezone.now().strftime('%Y-%m-%d %H:%M:%S')
-            # f.write("\r\n{} Task finish\r\n".format(date_end))
+            date_end = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+            print("\r\n{} Task finished".format(date_end))
             history.is_finished = True
             if summary.get('dark'):
                 history.is_success = False

apps/ops/templates/ops/task_list.html

@@ -113,7 +113,8 @@ $(document).ready(function() {
     };
     var success = function(data) {
         var task_id = data.task;
-        window.location = "{% url 'ops:adhoc-history-output' pk=DEFAULT_PK %}".replace("{{ DEFAULT_PK }}", task_id);
+        var url = '{% url "common:celery-task-log" pk=DEFAULT_PK %}'.replace("{{ DEFAULT_PK }}", task_id);
+        window.open(url, '', 'width=800,height=800')
     };
     APIUpdateAttr({
         url: the_url,