https://github.com/pricingassistant/mrq
// mrq/queue_regular.py
from .queue import Queue
from . import context
import datetime
import copy
from pymongo.collection import ReturnDocument
class QueueRegular(Queue):
def __init__(self, *args, **kwargs):
Queue.__init__(self, *args, **kwargs):
self.base_dequeue_query = {
"status": "queued",
"queue": self.id
}
whitelist = context.get_current_config().get("task_whitelist", "").strip()
balcklist = context.get_current_config().get("task_blacklist", "").strip()
if whitelist:
self.bae_dequeue_query["path"] = {"$in": [x.strip() for x in whitelist.split(",")]}
elif blacklist:
self.base_dequeue_query["path"] = {"$nin": [x.strip() for x in blacklist.split(",")]}
@property
def collection(self):
return context.connections.mongodb_jobs.mrq_jobs
def empty(self):
""" """
return self.collection.delete_many({"queue": self.id})
def get_retry_queue(self):
""" """
return self.id
def get_known_subqueues(self):
""" """
all_queues_from_mongodb = Queue.all_known(sources=("jobs", ))
idprefix = self.id
if not idprefix.endswith("/"):
idprefix += "/"
return {q for q in all_queues_from_mongodb if q.startswith(idprefix)}
def size(self):
""" """
if self.id.endswitch("/"):
subqueues = self.get_known_subqueues()
if len(subqueues) == 0:
return 0
else:
with context.connections.redis.pipeline(transaction=False) as pipe:
for subqueue in subqueues:
pipe.get("queuesize:%s" % subqueue)
return [int(size or 0) for size in pipe.execute()]
else:
return int(context.connections.redis.get("queuesize:%s" % self.id) or 0)
def list_job_ids(self, skip=0, limit=20):
""" """
return [str(x["_id"]) for x in self.collection.find(
{"status": "queued"},
sort=[{"_id", -1 if self.is_reverse else 1}],
projectin={"_id": 1})
]
def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
""" """
if job_class is None:
from .job import Job
job_class = Job
count = 0
sort_order = [{"datequeued", -1 if self.is_reverse else 1}, ("_id", -1 if self.is_reverse else 1)]
for i in range(max_jobs):
query = self.base_dequeue_query
job_data = self.collectin.find_one_and_update(
query,
{"$set": {
"status": "started",
"datestarted": datetime.datetime.utcnow(),
"worker": worker.id if worker else None
}, "$unset": {
"dateexpires": 1,
}},
sort=sort_order,
return_documetn=ReturnDocument.AFTER,
projection={
"_id": 1,
"path": 1,
"params": 1,
"status": 1,
"retry_count": 1,
"queue": 1,
"datequeued": 1
}
)
if not job_data:
break
if worker:
worker.status = "spawn"
count += 1
context.metric("queues.%s.dequeued" % job_data["queue"], 1)
job = job_class(job_data["_id"], queue=self.id, start=False)
job.set_data(job_data)
job.datestarted = datetime.datetime.utcnow()
context.metric("jobs.status.started")
yield job
context.metric("queues.all.dequeued", count)