From d3283571a6b15d52a03b652799cf4cfded7ac308 Mon Sep 17 00:00:00 2001
From: Ilan Steemers <koed00@gmail.com>
Date: Sun, 30 Aug 2015 14:33:35 +0200
Subject: [PATCH] First draft of a Disque broker

TODO monitoring doesn't work yet
---
 django_q/brokers/__init__.py | 12 ++++++++----
 django_q/brokers/disque.py   | 36 ++++++++++++++++++++++++++++++++++++
 django_q/conf.py             |  8 +++++++-
 3 files changed, 51 insertions(+), 5 deletions(-)
 create mode 100644 django_q/brokers/disque.py

diff --git a/django_q/brokers/__init__.py b/django_q/brokers/__init__.py
index 1d6f746..213ed49 100644
--- a/django_q/brokers/__init__.py
+++ b/django_q/brokers/__init__.py
@@ -39,9 +39,13 @@ class Broker(object):
 
 
 def get_broker(list_key=Conf.Q_LIST):
-    if Conf.REDIS:
-        from brokers import redis
-        return redis.Redis(list_key=list_key)
-    elif Conf.DJANGO_REDIS:
+    if Conf.DJANGO_REDIS:
         from brokers import djangoredis
         return djangoredis.DjangoRedis(list_key=list_key)
+    elif Conf.DISQUE:
+        from brokers import disque
+        return disque.Disque(list_key=list_key)
+    # default to redis
+    else:
+        from brokers import redis
+        return redis.Redis(list_key=list_key)
diff --git a/django_q/brokers/disque.py b/django_q/brokers/disque.py
new file mode 100644
index 0000000..faff2c9
--- /dev/null
+++ b/django_q/brokers/disque.py
@@ -0,0 +1,36 @@
+import redis
+from django_q.brokers import Broker
+from django_q.conf import Conf
+
+
+class Disque(Broker):
+    def enqueue(self, task):
+        return self.connection.execute_command(
+            'ADDJOB {} {} 500 RETRY {}'.format(self.list_key, task, Conf.RETRY)).decode()
+
+    def dequeue(self):
+        task = self.connection.execute_command('GETJOB TIMEOUT 1000 FROM {}'.format(self.list_key))
+        if task:
+            return task[0][1].decode(), task[0][2].decode()
+
+    def queue_size(self):
+        return self.connection.execute_command('QLEN {}'.format(self.list_key))
+
+    def acknowledge(self, ack_id):
+        return self.connection.execute_command('ACKJOB {}'.format(ack_id))
+
+    def ping(self):
+        return self.connection.ping()
+
+    @staticmethod
+    def get_connection():
+        for node in Conf.DISQUE:
+            host, port = node.split(':')
+            redis_client = redis.Redis(host, int(port))
+            try:
+                redis_client.ping()
+                redis_client.decode_responses = True
+                return redis_client
+            except redis.exceptions.ConnectionError:
+                pass
+        raise ConnectionError('Could not connect to any Disque nodes')
diff --git a/django_q/conf.py b/django_q/conf.py
index c31834c..615abf3 100644
--- a/django_q/conf.py
+++ b/django_q/conf.py
@@ -33,6 +33,9 @@ class Conf(object):
 
     DJANGO_REDIS = conf.get('django_redis', None)
 
+    # Disque broker
+    DISQUE = conf.get('disque', None)
+
     # Name of the cluster or site. For when you run multiple sites on one redis server
     PREFIX = conf.get('name', 'default')
 
@@ -46,7 +49,6 @@ class Conf(object):
     # Maximum number of tasks that each cluster can work on
     QUEUE_LIMIT = conf.get('queue_limit', None)
 
-
     # Number of workers in the pool. Default is cpu count if implemented, otherwise 4.
     WORKERS = conf.get('workers', False)
     if not WORKERS:
@@ -70,6 +72,10 @@ class Conf(object):
     # Number of seconds to wait for a worker to finish.
     TIMEOUT = conf.get('timeout', None)
 
+    # Number of seconds to wait for acknowledgement before retrying a task
+    # Only works with brokers that guarantee delivery. Defaults to 0. Meaning no retries.
+    RETRY = conf.get('retry', 0)
+
     # The Django Admin label for this app
     LABEL = conf.get('label', 'Django Q')