django-q2/django_q/brokers/orm.py

from datetime import timedelta
from time import sleep
from django.utils import timezone
from django.db.models import Q
from django_q.brokers import Broker
from django_q.models import OrmQ
from django_q.conf import Conf


class ORM(Broker):
    def queue_size(self):
        return OrmQ.objects.using(Conf.ORM) \
            .filter(Q(key=self.list_key, lock__isnull=True) |
                    Q(key=self.list_key, lock__lte=timezone.now() - timedelta(seconds=Conf.RETRY))) \
            .count()

    def purge_queue(self):
        return OrmQ.objects.using(Conf.ORM).filter(key=self.list_key).delete()

    def ping(self):
        return True

    def info(self):
        return 'ORM {}'.format(Conf.ORM)

    def fail(self, task_id):
        self.delete(task_id)

    def enqueue(self, task):
        package = OrmQ.objects.using(Conf.ORM).create(key=self.list_key, payload=task)
        return package.pk

    def dequeue(self):
        if len(self.task_cache) > 0:
            t = self.task_cache.pop()
            return t.pk, t.payload
        else:
            # Get new and timed out tasks
            tasks = OrmQ.objects.using(Conf.ORM).filter(
                Q(key=self.list_key, lock__isnull=True) |
                Q(key=self.list_key, lock__lte=timezone.now() - timedelta(seconds=Conf.RETRY)))[:Conf.BULK]
            if tasks:
                # lock them
                OrmQ.objects.using(Conf.ORM).filter(pk__in=tasks).update(lock=timezone.now())
                tasks = [t for t in tasks]
                # pop one task
                t = tasks.pop()
                if tasks:
                    # add remainder to cache
                    self.task_cache = [t for t in tasks]
                return t.pk, t.payload
            # empty queue, spare the cpu
            sleep(0.2)

    def delete_queue(self):
        return self.purge_queue()

    def delete(self, task_id):
        return OrmQ.objects.using(Conf.ORM).filter(pk=task_id).delete()

    def acknowledge(self, task_id):
        return self.delete(task_id)