openedx / django-user-tasks Goto Github PK

from celery import signals
from user_tasks.signals import create_user_task
from myproject.celery.celery import app

signals.before_task_publish.disconnect(create_user_task)
signals.before_task_publish.connect(create_user_task_wrapper)

def create_user_task_wrapper(sender=None, body=None, **kwargs):
    return create_user_task(
        sender,
        body
        if app.conf.task_protocol == 1
        else proto2_to_proto1(body, kwargs.get("headers", {})),
    )

# Based on some old celery code:
def proto2_to_proto1(body, headers):
    args, kwargs, embed = body
    embedded = _extract_proto2_embed(**embed)
    chained = embedded.pop("chain")
    new_body = dict(
        _extract_proto2_headers(**headers), args=args, kwargs=kwargs, **embedded
    )
    if chained:
        new_body["callbacks"].append(chain(chained))
    return new_body


def _extract_proto2_headers(id, retries, eta, expires, group, timelimit, task, **_):
    return {
        "id": id,
        "task": task,
        "retries": retries,
        "eta": eta,
        "expires": expires,
        "utc": True,
        "taskset": group,
        "timelimit": timelimit,
    }


def _extract_proto2_embed(callbacks, errbacks, chain, chord, **_):
    return {
        "callbacks": callbacks or [],
        "errbacks": errbacks,
        "chain": chain,
        "chord": chord,
    }

class ImportTask(UserTask):
    @classmethod
    def generate_name(cls, arguments_dict):
        serie = Serie.objects.get(pk=arguments_dict['serie_id'])
        return 'Import of {}'.format(serie)

    @staticmethod
    def calculate_total_steps(arguments_dict):
        return 100


@shared_task(base=ImportTask, bind=True)
def import_csv(self, user_id, notification_key, path_file, serie_id, **kwargs):
    uploadCandles(path_file, serie_id)

def user_logged_in_handler(sender, request, user, **kwargs):
    print('OUI')


def user_logged_out_handler(sender, request, user, **kwargs):
    print('NON')


def user_task_stopped_handler(**kwargs):
    print('User task stoped')


@receiver(task_success)
def task_success_handler(result, **kwargs):
    print('User task succeed')


@task_success.connect
def task_succeeded_handler(sender=None, **kwargs):
    print(sender)
    if isinstance(sender, UserTaskMixin):
        print('User task succeed')


user_logged_in.connect(user_logged_in_handler)
user_logged_out.connect(user_logged_out_handler)
user_task_stopped.connect(user_task_stopped_handler)