[Update] 修改 task 运行机制

apps/ops/utils.py

@@ -1,123 +1,36 @@
 # ~*~ coding: utf-8 ~*~
+import json
+from django_celery_beat.models import PeriodicTask, IntervalSchedule, CrontabSchedule
 
-import time
-from django.utils import timezone
-from django.db import transaction
-from django_celery_beat.models import PeriodicTask, IntervalSchedule
 
-from common.utils import get_logger, get_object_or_none, get_short_uuid_str
-from .ansible import AdHocRunner, CommandResultCallback
-from .inventory import JMSInventory
-from .ansible.exceptions import AnsibleError
-from .models import AdHocRunHistory, Task, AdHoc
+from common.utils import get_logger, get_object_or_none
+from .models import Task, AdHoc
 
 logger = get_logger(__file__)
 
 
-def record_adhoc(func):
-    def _deco(adhoc, **options):
-        record = AdHocRunHistory(adhoc=adhoc, task=adhoc.task)
-        time_start = time.time()
-        try:
-            result = func(adhoc, **options)
-            record.is_finished = True
-            if result.results_summary.get('dark'):
-                record.is_success = False
-            else:
-                record.is_success = True
-            record.result = result.results_raw
-            record.summary = result.results_summary
-            return result
-        finally:
-            record.date_finished = timezone.now()
-            record.timedelta = time.time() - time_start
-            record.save()
-    return _deco
+def get_task_by_id(task_id):
+    return get_object_or_none(Task, id=task_id)
 
 
-def get_adhoc_inventory(adhoc):
-    if adhoc.become:
-        become_info = {
-            'become': {
-               adhoc.become
-            }
-        }
-    else:
-        become_info = None
-
-    inventory = JMSInventory(
-        adhoc.hosts, run_as_admin=adhoc.run_as_admin,
-        run_as=adhoc.run_as, become_info=become_info
-    )
-    return inventory
-
-
-def get_inventory(hostname_list, run_as_admin=False, run_as=None, become_info=None):
-    return JMSInventory(
-        hostname_list, run_as_admin=run_as_admin,
-        run_as=run_as, become_info=become_info
-    )
-
-
-def get_adhoc_runner(hostname_list, run_as_admin=False, run_as=None, become_info=None):
-    inventory = get_inventory(
-        hostname_list, run_as_admin=run_as_admin,
-        run_as=run_as, become_info=become_info
-    )
-    runner = AdHocRunner(inventory)
-    return runner
-
-
-@record_adhoc
-def run_adhoc_object(adhoc, **options):
-    """
-    :param adhoc: Instance of AdHoc
-    :param options: ansible support option, like forks ...
-    :return:
-    """
-    name = adhoc.task.name
-    inventory = get_adhoc_inventory(adhoc)
-    runner = AdHocRunner(inventory)
-    for k, v in options.items():
-        runner.set_option(k, v)
-
-    try:
-        result = runner.run(adhoc.tasks, adhoc.pattern, name)
-        return result
-    except AnsibleError as e:
-        logger.error("Failed run adhoc {}, {}".format(name, e))
-        raise
-
-
-def run_adhoc(hostname_list, pattern, tasks, name=None,
-              run_as_admin=False, run_as=None, become_info=None):
-    if name is None:
-        name = "Adhoc-task-{}-{}".format(
-            get_short_uuid_str(),
-            timezone.now().strftime("%Y-%m-%d %H:%M:%S"),
-        )
-
-    inventory = get_inventory(
-        hostname_list, run_as_admin=run_as_admin,
-        run_as=run_as, become_info=become_info
-    )
-    runner = AdHocRunner(inventory)
-    return runner.run(tasks, pattern, play_name=name)
-
-
-def create_or_update_task(
+def create_or_update_ansible_task(
         task_name, hosts, tasks, pattern='all', options=None,
         run_as_admin=False, run_as="", become_info=None,
-        created_by=None
+        created_by=None, interval=None, crontab=None,
+        is_periodic=False, callback=None,
     ):
-    print(options)
-    print(task_name)
+
     task = get_object_or_none(Task, name=task_name)
+
     if task is None:
-        task = Task(name=task_name, created_by=created_by)
+        task = Task(
+            name=task_name, interval=interval,
+            crontab=crontab, is_periodic=is_periodic,
+            callback=callback, created_by=created_by
+        )
         task.save()
 
-    adhoc = task.get_latest_adhoc()
+    adhoc = task.latest_adhoc
     new_adhoc = AdHoc(task=task, pattern=pattern,
                       run_as_admin=run_as_admin,
                       run_as=run_as)
@@ -128,23 +41,67 @@ def create_or_update_task(
     if not adhoc or adhoc != new_adhoc:
         new_adhoc.save()
         task.latest_adhoc = new_adhoc
-    print("Return task")
     return task
 
 
-def create_periodic_tasks(tasks):
+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():
-        schedule, _ = IntervalSchedule.objects.get_or_create(
-            every=detail['schedule'],
-            period=IntervalSchedule.SECONDS,
-        )
+        interval = None
+        crontab = None
+        if isinstance(detail.get("interval"), int):
+            interval, _ = IntervalSchedule.objects.get_or_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")
 
-        task = PeriodicTask.objects.create(
-            interval=schedule,
+            crontab, _ = CrontabSchedule.objects.get_or_create(
+                minute=minute, hour=hour, day_of_week=week,
+                day_of_month=day, month_of_year=month,
+            )
+        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['enabled']
         )
-        print("Create periodic task: {}".format(task))
+
+        task = PeriodicTask.objects.update_or_create(
+            defaults=defaults, name=name,
+        )
+        logger.info("Create periodic task: {}".format(task))
+        return task
+
+
+def disable_celery_periodic_task(task_name):
+    PeriodicTask.objects.filter(name=task_name).update(enabled=False)
+
+
+def delete_celery_periodic_task(task_name):
+    PeriodicTask.objects.filter(name=task_name).delete()
+