Django ORM broker

django_q/brokers/orm.py

@@ -0,0 +1,62 @@
+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)