diff --git a/mrq/config.py b/mrq/config.py index 5ee51425..ad772fab 100644 --- a/mrq/config.py +++ b/mrq/config.py @@ -3,7 +3,7 @@ import sys import re from .version import VERSION -from .utils import get_local_ip +from .utils import get_local_ip, DelimiterArgParser import atexit @@ -344,7 +344,8 @@ def add_parser_args(parser, config_type): parser.add_argument( '--subqueues_delimiter', default='/', - help='Delimiter between main queue and subqueue names') + help='Delimiter between main queue and subqueue names', + action=DelimiterArgParser) parser.add_argument( '--admin_port', diff --git a/mrq/queue.py b/mrq/queue.py index 115238d2..3281af97 100644 --- a/mrq/queue.py +++ b/mrq/queue.py @@ -15,6 +15,11 @@ class Queue(object): is_set = False is_reverse = False + # root_id will contain the root queue id without any trailing subqueue delimiter + # e.g. if self.id is "some_queue/" then self.root_id will contain "some_queue" + # and if self.id is "some_queue/some_subqueue" then self.root_id will contain "some_queue" + root_id = None + use_large_ids = False # This is a mutable type so it is shared by all instances @@ -47,6 +52,13 @@ def __init__(self, queue_id, add_to_known_queues=False): if "_sorted" in self.id: self.is_sorted = True + self.root_id = self.id + + delimiter = context.get_current_config().get("subqueues_delimiter") + if delimiter is not None and delimiter in self.id: + # Get the root queue id with no trailing delimiter + self.root_id = self.id.split(delimiter)[0] + self.use_large_ids = context.get_current_config()["use_large_job_ids"] # If this is the first time this process sees this queue, try to add it @@ -117,7 +129,7 @@ def redis_known_subqueues(self): def get_config(self): """ Returns the specific configuration for this queue """ - return context.get_current_config().get("raw_queues", {}).get(self.id) or {} + return context.get_current_config().get("raw_queues", {}).get(self.root_id) or {} def serialize_job_ids(self, job_ids): """ Returns job_ids serialized for storage in Redis """ diff --git a/mrq/utils.py b/mrq/utils.py index 8af6a455..734f1dff 100644 --- a/mrq/utils.py +++ b/mrq/utils.py @@ -4,6 +4,7 @@ import math import json import datetime +import argparse from collections import deque from bson import ObjectId @@ -194,3 +195,11 @@ def next(self, val): self.__sum += val self.__q.append(val) return 1.0 * self.__sum / len(self.__q) + + +class DelimiterArgParser(argparse.Action): + def __call__(self, parser, namespace, value, option_string): + if value == '_': + parser.error("Cannot use '%s' as a subqueue delimiter" % value) + + setattr(namespace, self.dest, value) diff --git a/tests/test_raw.py b/tests/test_raw.py index fb1f1f32..cefb2c45 100644 --- a/tests/test_raw.py +++ b/tests/test_raw.py @@ -79,14 +79,23 @@ def test_raw_sorted(worker, p_queue, p_pushback, p_timed, p_flags): assert test_collection.count() == 3 +@pytest.mark.parametrize("has_subqueue", [False, True]) @pytest.mark.parametrize(["p_queue", "p_set"], [ ["test_raw", False], ["test_set", True] ]) -def test_raw_set(worker, p_queue, p_set): - - worker.start( - flags="--greenlets 10 --config tests/fixtures/config-raw1.py", queues=p_queue) +def test_raw_set(worker, has_subqueue, p_queue, p_set): + flags = "--greenlets 10 --config tests/fixtures/config-raw1.py" + if has_subqueue: + flags = "%s --subqueues_refresh_interval=0.1" % flags + # worker should dequeue all subqueues + p_queue = "%s/" % p_queue + + worker.start(flags=flags, queues=p_queue) + + if has_subqueue: + # queue tasks in p_queue/subqueue + p_queue = "%ssubqueue" % p_queue test_collection = worker.mongodb_logs.tests_inserts jobs_collection = worker.mongodb_jobs.mrq_jobs diff --git a/tests/test_subqueues.py b/tests/test_subqueues.py index 27f5cea7..34b65f33 100644 --- a/tests/test_subqueues.py +++ b/tests/test_subqueues.py @@ -43,7 +43,7 @@ def test_unmatchable_subqueues(worker, queue, enqueue_on): worker.stop() -@pytest.mark.parametrize(["delimiter"], ["/", ".", "_"]) +@pytest.mark.parametrize(["delimiter"], ["/", ".", "-"]) def test_custom_delimiters(worker, delimiter): queue = "main" + delimiter