From 191cc2854348bc6799553413c35ed0d8c0ec0ab1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BB=BB=E6=99=93=E7=A3=8A?= Date: Fri, 12 Aug 2016 17:07:30 +0800 Subject: [PATCH] custom connection class --- rq/cli/cli.py | 12 +++++++----- rq/cli/helpers.py | 6 +++--- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/rq/cli/cli.py b/rq/cli/cli.py index bb3c5a6..0b3e611 100755 --- a/rq/cli/cli.py +++ b/rq/cli/cli.py @@ -37,12 +37,12 @@ config_option = click.option('--config', '-c', help='Module containing RQ settings.') -def connect(url, config=None): +def connect(url, config=None, connection_class='redis.StrictRedis'): if url: - return StrictRedis.from_url(url) + return connection_class.from_url(url) settings = read_config_file(config) if config else {} - return get_redis_from_config(settings) + return get_redis_from_config(settings, connection_class) @click.group() @@ -143,6 +143,7 @@ def info(url, config, path, interval, raw, only_queues, only_workers, by_queue, @click.option('--worker-class', '-w', default='rq.Worker', help='RQ Worker class to use') @click.option('--job-class', '-j', default='rq.job.Job', help='RQ Job class to use') @click.option('--queue-class', default='rq.Queue', help='RQ Queue class to use') +@click.option('--connection-class', default='redis.StrictRedis', help='Redis client class to use') @click.option('--path', '-P', default='.', help='Specify the import path.') @click.option('--results-ttl', type=int, help='Default results timeout to be used') @click.option('--worker-ttl', type=int, help='Default worker timeout to be used') @@ -152,7 +153,7 @@ def info(url, config, path, interval, raw, only_queues, only_workers, by_queue, @click.option('--exception-handler', help='Exception handler(s) to use', multiple=True) @click.option('--pid', help='Write the process ID number to a file at the specified path') @click.argument('queues', nargs=-1) -def worker(url, config, burst, name, worker_class, job_class, queue_class, path, results_ttl, worker_ttl, +def worker(url, config, burst, name, worker_class, job_class, queue_class, connection_class, path, results_ttl, worker_ttl, verbose, quiet, sentry_dsn, exception_handler, pid, queues): """Starts an RQ worker.""" @@ -170,7 +171,8 @@ def worker(url, config, burst, name, worker_class, job_class, queue_class, path, setup_loghandlers_from_args(verbose, quiet) - conn = connect(url, config) + connection_class = import_attribute(connection_class) + conn = connect(url, config, connection_class) cleanup_ghosts(conn) worker_class = import_attribute(worker_class) queue_class = import_attribute(queue_class) diff --git a/rq/cli/helpers.py b/rq/cli/helpers.py index 7bfc0cf..340233d 100644 --- a/rq/cli/helpers.py +++ b/rq/cli/helpers.py @@ -26,10 +26,10 @@ def read_config_file(module): if k.upper() == k]) -def get_redis_from_config(settings): +def get_redis_from_config(settings, connection_class): """Returns a StrictRedis instance from a dictionary of settings.""" if settings.get('REDIS_URL') is not None: - return StrictRedis.from_url(settings['REDIS_URL']) + return connection_class.from_url(settings['REDIS_URL']) kwargs = { 'host': settings.get('REDIS_HOST', 'localhost'), @@ -52,7 +52,7 @@ def get_redis_from_config(settings): if not version_info >= (2, 10): raise RuntimeError('Using SSL requires a redis-py version >= 2.10') kwargs['ssl'] = use_ssl - return StrictRedis(**kwargs) + return connection_class(**kwargs) def pad(s, pad_to_length):