Give queue name as second raw job factory argument when requested

mrq/queue_raw.py

@@ -195,7 +195,7 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
         if worker:
             worker.status = "spawn"
 
-        job_data = [job_factory(p) for p in params]
+        job_data = [job_factory(*p) for p in self.get_job_factory_params(job_factory, params)]
         for j in job_data:
             j["status"] = "started"
             j["queue"] = retry_queue
@@ -206,6 +206,22 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
         for job in job_class.insert(job_data, statuses_no_storage=statuses_no_storage):
             yield job
 
+    def get_job_factory_params(self, job_factory, params):
+        args_count = job_factory.__code__.co_argcount
+
+        if args_count == 0:
+            raise Exception("Job factory must take at least 1 argument!")
+
+        if args_count == 1:
+            # just expecting raw params
+            return ([p] for p in params)
+
+        if args_count == 2:
+            # expecting raw queue name as second argument
+            return ([p, self.id] for p in params)
+
+        raise Exception("Job factory expecting too many arguments!")
+
     def get_sorted_graph(
             self,
             start=0,

tests/fixtures/config-raw1.py

@@ -147,5 +147,14 @@
                 "sleep": float(rawparam.split(" ")[1])
             }
         }
+    },
+    "test_factory_params_raw": {
+        "job_factory": lambda rawparam, queue: {
+            "path": "tests.tasks.general.MongoInsert",
+            "params": {
+                "queue": queue,
+                "sleep": float(rawparam)
+            }
+        }
     }
 }

tests/test_raw.py

@@ -360,3 +360,26 @@ def test_raw_no_storage(worker):
     time.sleep(2)
 
     assert jobs_collection.count({"status": "retry"}) == 1
+
+
+def test_raw_factory_params(worker):
+
+    p_queue = "test_factory_params_raw"
+
+    worker.start(
+        flags="--greenlets 10 --config tests/fixtures/config-raw1.py", queues=p_queue)
+
+    jobs_collection = worker.mongodb_jobs.mrq_jobs
+    tests_collection = worker.mongodb_jobs.tests_inserts
+
+    assert jobs_collection.count() == 0
+    assert tests_collection.count() == 0
+    assert Queue(p_queue).size() == 0
+
+    worker.send_raw_tasks(p_queue, [0], block=True)
+
+    assert Queue(p_queue).size() == 0
+    assert jobs_collection.count() == 1
+    assert jobs_collection.count({"status": "success"}) == 1
+    assert tests_collection.count() == 1
+    assert tests_collection.count({"params.queue": p_queue, "params.sleep": 0.0}) == 1