Adds a MongoDB broker

* adds mongo broker tests
* adds mongo brokers docs

* broker info is now cached when it needs to connect to a server, to reduce traffic.

django_q/brokers/__init__.py

@@ -7,6 +7,7 @@ class Broker(object):
         self.connection = self.get_connection(list_key)
         self.list_key = list_key
         self.cache = self.get_cache()
+        self._info = None
 
     def enqueue(self, task):
         """
@@ -78,7 +79,7 @@ class Broker(object):
         """
         Shows the broker type
         """
-        pass
+        return self._info
 
     def set_stat(self, key, value, timeout):
         """
@@ -167,6 +168,9 @@ def get_broker(list_key=Conf.PREFIX):
     elif Conf.ORM:
         from brokers import orm
         return orm.ORM(list_key=list_key)
+    elif Conf.MONGO:
+        from brokers import mongo
+        return mongo.Mongo(list_key=list_key)
     # default to redis
     else:
         from brokers import redis_broker

django_q/brokers/disque.py

@@ -40,8 +40,10 @@ class Disque(Broker):
         return len(jobs)
 
     def info(self):
-        info = self.connection.info('server')
-        return 'Disque {}'.format(info['disque_version'])
+        if not self._info:
+            info = self.connection.info('server')
+            self._info= 'Disque {}'.format(info['disque_version'])
+        return self._info
 
     @staticmethod
     def get_connection(list_key=Conf.PREFIX):

django_q/brokers/mongo.py

@@ -0,0 +1,64 @@
+from datetime import timedelta
+from time import sleep
+from bson import ObjectId
+
+from django.utils import timezone
+
+from pymongo import MongoClient
+
+from django_q.brokers import Broker
+from django_q.conf import Conf
+
+
+def _timeout():
+    return timezone.now() - timedelta(seconds=Conf.RETRY)
+
+
+class Mongo(Broker):
+    def __init__(self, list_key=Conf.PREFIX):
+        super(Mongo, self).__init__(list_key)
+        self.collection = self.connection[Conf.MONGO_DB][list_key]
+
+    @staticmethod
+    def get_connection(list_key=Conf.PREFIX):
+        return MongoClient(**Conf.MONGO)
+
+    def queue_size(self):
+        return self.collection.count({'lock': {'$lte': _timeout()}})
+
+    def lock_size(self):
+        return self.collection.count({'lock': {'$gt': _timeout()}})
+
+    def purge_queue(self):
+        return self.delete_queue()
+
+    def ping(self):
+        return self.info is not None
+
+    def info(self):
+        if not self._info:
+            self._info = 'MongoDB {}'.format(self.connection.server_info()['version'])
+        return self._info
+
+    def fail(self, task_id):
+        self.delete(task_id)
+
+    def enqueue(self, task):
+        inserted_id = self.collection.insert_one({'payload': task, 'lock': _timeout()}).inserted_id
+        return str(inserted_id)
+
+    def dequeue(self):
+        task = self.collection.find_one_and_update({'lock': {'$lte': _timeout()}}, {'$set': {'lock': timezone.now()}})
+        if task:
+            return [(str(task['_id']), task['payload'])]
+        # empty queue, spare the cpu
+        sleep(0.2)
+
+    def delete_queue(self):
+        return self.collection.drop()
+
+    def delete(self, task_id):
+        self.collection.delete_one({'_id': ObjectId(task_id)})
+
+    def acknowledge(self, task_id):
+        return self.delete(task_id)

django_q/brokers/orm.py

@@ -21,7 +21,7 @@ class ORM(Broker):
         return self.connection.filter(key=self.list_key, lock__lte=_timeout()).count()
 
     def lock_size(self):
-        return self.connection.filter(key=self.list_key, lock__gte=_timeout()).count()
+        return self.connection.filter(key=self.list_key, lock__gt=_timeout()).count()
 
     def purge_queue(self):
         return self.connection.filter(key=self.list_key).delete()
@@ -30,7 +30,9 @@ class ORM(Broker):
         return True
 
     def info(self):
-        return 'ORM {}'.format(Conf.ORM)
+        if not self._info:
+            self._info = 'ORM {}'.format(Conf.ORM)
+        return self._info
 
     def fail(self, task_id):
         self.delete(task_id)

django_q/brokers/redis_broker.py

@@ -1,4 +1,5 @@
 import redis
+
 from django_q.brokers import Broker
 from django_q.conf import Conf, logger
 
@@ -9,7 +10,6 @@ except ImportError:
 
 
 class Redis(Broker):
-
     def __init__(self, list_key=Conf.PREFIX):
         super(Redis, self).__init__(list_key='django_q:{}:q'.format(list_key))
 
@@ -38,8 +38,10 @@ class Redis(Broker):
             raise e
 
     def info(self):
-        info = self.connection.info('server')
-        return 'Redis {}'.format(info['redis_version'])
+        if not self._info:
+            info = self.connection.info('server')
+            self._info = 'Redis {}'.format(info['redis_version'])
+        return self._info
 
     def set_stat(self, key, value, timeout):
         self.connection.set(key, value, timeout)

django_q/conf.py

@@ -50,6 +50,10 @@ class Conf(object):
     # ORM broker
     ORM = conf.get('orm', None)
 
+    # MongoDB broker
+    MONGO = conf.get('mongo', None)
+    MONGO_DB = conf.get('mongo_db', 'django-q')
+
     # Name of the cluster or site. For when you run multiple sites on one redis server
     PREFIX = conf.get('name', 'default')
 

django_q/monitor.py

@@ -86,7 +86,7 @@ def monitor(run_once=False, broker=None):
             lock_size = broker.lock_size()
             if lock_size:
                 queue_size = '{}({})'.format(queue_size, lock_size)
-            print(term.move(i, 0) + term.white_on_cyan(term.center(broker.info(), width=col_width * 2)))
+            print(term.move(i, 0) + term.white_on_cyan(term.center(broker.info, width=col_width * 2)))
             print(term.move(i, 2 * col_width) + term.black_on_cyan(term.center(_('Queued'), width=col_width)))
             print(term.move(i, 3 * col_width) + term.white_on_cyan(term.center(queue_size, width=col_width)))
             print(term.move(i, 4 * col_width) + term.black_on_cyan(term.center(_('Success'), width=col_width)))

django_q/tests/test_brokers.py

@@ -285,3 +285,64 @@ def test_orm():
     assert broker.queue_size() == 0
     # back to django-redis
     Conf.ORM = None
+
+
+@pytest.mark.django_db
+def test_mongo():
+    Conf.MONGO = {'host': '127.0.0.1', 'port': 27017}
+    # check broker
+    broker = get_broker(list_key='mongo_test')
+    assert broker.ping() is True
+    assert broker.info() is not None
+    # clear before we start
+    broker.delete_queue()
+    # enqueue
+    broker.enqueue('test')
+    assert broker.queue_size() == 1
+    # dequeue
+    task = broker.dequeue()[0]
+    assert task[1] == 'test'
+    broker.acknowledge(task[0])
+    assert broker.queue_size() == 0
+    # Retry test
+    Conf.RETRY = 1
+    broker.enqueue('test')
+    assert broker.queue_size() == 1
+    broker.dequeue()
+    assert broker.queue_size() == 0
+    sleep(1.5)
+    assert broker.queue_size() == 1
+    task = broker.dequeue()[0]
+    assert broker.queue_size() == 0
+    broker.acknowledge(task[0])
+    sleep(1.5)
+    assert broker.queue_size() == 0
+    # delete job
+    task_id = broker.enqueue('test')
+    broker.delete(task_id)
+    assert broker.dequeue() is None
+    # fail
+    task_id = broker.enqueue('test')
+    broker.fail(task_id)
+    # bulk test
+    for i in range(5):
+        broker.enqueue('test')
+    tasks = []
+    for i in range(5):
+        tasks.append(broker.dequeue()[0])
+    assert broker.lock_size() == 5
+    for task in tasks:
+        assert task is not None
+        broker.acknowledge(task[0])
+    # test lock size
+    assert broker.lock_size() == 0
+    # test duplicate acknowledge
+    broker.acknowledge(task[0])
+    # delete queue
+    broker.enqueue('test')
+    broker.enqueue('test')
+    broker.purge_queue()
+    broker.delete_queue()
+    assert broker.queue_size() == 0
+    # back to django-redis
+    Conf.ORM = None

django_q/tests/test_cluster.py

@@ -33,6 +33,7 @@ def broker():
     Conf.IRON_MQ = None
     Conf.SQS = None
     Conf.ORM = None
+    Conf.MONGO = None
     Conf.DJANGO_REDIS = 'default'
     return get_broker()