decorator.py 1.86 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
from ippidb.models import Job
from celery import Task, states
from django_celery_results.models import TaskResult


class AlreadyExistError(Exception):
    pass


class IppidbTask(Task):
    """
    Ippidb custom task
    """

    def __call__(self, *args, **kwargs):
        """In celery task this function call the run method, here you can
        set some environment variable before the run of the task"""

        tasks = TaskResult.objects.filter(
            task_name=self.request.task,
            status__in=[states.STARTED, states.PENDING, states.RETRY, states.RECEIVED],
        )
        count_tasks = tasks.count()
        if not count_tasks:
            self.update_state(state=states.PENDING)
            return self.run(*args, **kwargs)
        else:
            message_exc = "Job {} in state {} already exist".format(
                tasks[0].task_name, tasks[0].status
            )
            raise AlreadyExistError(message_exc)

    def update_job(self, std_out=None, std_err=None):
        job = Job.objects.get(task_result__task_id=self.task_id)
        if std_out:
            job.write(std_out, output="std_out")
        elif std_err:
            job.write(std_err, output="std_err")

    def on_success(self, retval, task_id, args, kwargs):
        self.update_job(std_out="SUCCESS")
        super(IppidbTask, self).on_success(retval, task_id, args, kwargs)

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        self.update_job(std_err=exc)
        super(IppidbTask, self).on_failure(exc, task_id, args, kwargs, einfo)

    def update_state(self, task_id=None, state=None, meta=None, **kwargs):
        self.state = state
        if task_id is None:
            self.task_id = self.request.id
        else:
            self.task_id = task_id
        super(IppidbTask, self).update_state(
            task_id=task_id, state=state, meta=meta, **kwargs
        )