update known queues when requeuing jobs

Author: Florian Perucki

mrq/basetasks/utils.py

@@ -134,6 +134,12 @@ def perform_action(self, action, query, destination_queue):
                         "_id": {"$in": jobs_by_queue[queue]}
                     }, {"$set": updates}, multi=True)
 
+                if destination_queue is None:
+                    Queue.ensure_known_queues(jobs_by_queue.iterkeys())
+
+            if destination_queue is not None:
+                Queue.ensure_known_queues([destination_queue])
+
         print(stats)
 
         return stats

mrq/queue.py

@@ -10,6 +10,7 @@
 
 import sys
 from future import standard_library
+from itertools import chain
 
 PY3 = sys.version_info > (3,)
 standard_library.install_aliases()
@@ -119,6 +120,13 @@ def get_retry_queue(self):
         """ Return the name of the queue where retried jobs will be queued """
         return self.id
 
+    @classmethod
+    def ensure_known_queues(cls, queues):
+        """ List all previously known queues """
+        now = time.time()
+        params = chain.from_iterable((now, queue) for queue in queues)
+        context.connections.redis.zadd(Queue.redis_key_known_queues(), *params)
+
     def add_to_known_queues(self, timestamp=None):
         """ Adds this queue to the shared list of known queues """
         now = timestamp or time.time()

tests/test_general.py

@@ -179,6 +179,33 @@ def test_known_queues_lifecycle(worker):
     # Still not there.
     assert set(Queue.redis_known_queues().keys()) == set(["default", "xtest", "test_timed_set"])
 
+    # Now we're going to test that the known queues are correctly updated when requeuing a job
+
+    # Queue the job again
+    send_task("tests.tasks.general.Add", {"a": 41, "b": 1, "sleep": 1}, queue="x")
+
+    worker.send_task("mrq.basetasks.cleaning.CleanKnownQueues", {}, block=True)
+    # Requeue it in a different queue
+    params = {
+        "action": "requeue",
+        "destination_queue": "x2"
+    }
+    worker.send_task("mrq.basetasks.utils.JobAction", params, block=True)
+
+    assert set(Queue.redis_known_queues().keys()) == set(["x", "x2", "default", "xtest", "test_timed_set"])
+
+    Queue("x2").empty()
+    assert set(Queue.redis_known_queues().keys()) == set(["default", "x", "xtest", "test_timed_set"])
+
+    # Requeue it in the same queue
+    params = {
+        "action": "requeue"
+    }
+    worker.send_task("mrq.basetasks.utils.JobAction", params, block=True)
+
+    # The queue should be back
+    assert set(Queue.redis_known_queues().keys()) == set(["default", "x", "x2", "xtest", "test_timed_set"])
+
 
 def test_general_exception_status(worker):