# -*- coding: utf-8 -*- from __future__ import (absolute_import, division, print_function, unicode_literals) import sys import importlib import time from functools import partial import click import redis from redis import Redis from redis.sentinel import Sentinel from rq.defaults import (DEFAULT_CONNECTION_CLASS, DEFAULT_JOB_CLASS, DEFAULT_QUEUE_CLASS, DEFAULT_WORKER_CLASS) from rq.logutils import setup_loghandlers from rq.utils import import_attribute from rq.worker import WorkerStatus red = partial(click.style, fg='red') green = partial(click.style, fg='green') yellow = partial(click.style, fg='yellow') def read_config_file(module): """Reads all UPPERCASE variables defined in the given module file.""" settings = importlib.import_module(module) return dict([(k, v) for k, v in settings.__dict__.items() if k.upper() == k]) def get_redis_from_config(settings, connection_class=Redis): """Returns a StrictRedis instance from a dictionary of settings. To use redis sentinel, you must specify a dictionary in the configuration file. Example of a dictionary with keys without values: SENTINEL: {'INSTANCES':, 'SOCKET_TIMEOUT':, 'PASSWORD':,'DB':, 'MASTER_NAME':} """ if settings.get('REDIS_URL') is not None: return connection_class.from_url(settings['REDIS_URL']) elif settings.get('SENTINEL') is not None: instances = settings['SENTINEL'].get('INSTANCES', [('localhost', 26379)]) socket_timeout = settings['SENTINEL'].get('SOCKET_TIMEOUT', None) password = settings['SENTINEL'].get('PASSWORD', None) db = settings['SENTINEL'].get('DB', 0) master_name = settings['SENTINEL'].get('MASTER_NAME', 'mymaster') sn = Sentinel(instances, socket_timeout=socket_timeout, password=password, db=db) return sn.master_for(master_name) kwargs = { 'host': settings.get('REDIS_HOST', 'localhost'), 'port': settings.get('REDIS_PORT', 6379), 'db': settings.get('REDIS_DB', 0), 'password': settings.get('REDIS_PASSWORD', None), 'ssl': settings.get('REDIS_SSL', False), } return connection_class(**kwargs) def pad(s, pad_to_length): """Pads the given string to the given length.""" return ('%-' + '%ds' % pad_to_length) % (s,) def get_scale(x): """Finds the lowest scale where x <= scale.""" scales = [20, 50, 100, 200, 400, 600, 800, 1000] for scale in scales: if x <= scale: return scale return x def state_symbol(state): symbols = { WorkerStatus.BUSY: red('busy'), WorkerStatus.IDLE: green('idle'), WorkerStatus.SUSPENDED: yellow('suspended'), } try: return symbols[state] except KeyError: return state def show_queues(queues, raw, by_queue, queue_class, worker_class): num_jobs = 0 termwidth, _ = click.get_terminal_size() chartwidth = min(20, termwidth - 20) max_count = 0 counts = dict() for q in queues: count = q.count counts[q] = count max_count = max(max_count, count) scale = get_scale(max_count) ratio = chartwidth * 1.0 / scale for q in queues: count = counts[q] if not raw: chart = green('|' + '█' * int(ratio * count)) line = '%-12s %s %d' % (q.name, chart, count) else: line = 'queue %s %d' % (q.name, count) click.echo(line) num_jobs += count # print summary when not in raw mode if not raw: click.echo('%d queues, %d jobs total' % (len(queues), num_jobs)) def show_workers(queues, raw, by_queue, queue_class, worker_class): workers = set() for queue in queues: for worker in worker_class.all(queue=queue): workers.add(worker) if not by_queue: for worker in workers: queue_names = ', '.join(worker.queue_names()) name = '%s (%s %s)' % (worker.name, worker.hostname, worker.pid) if not raw: click.echo('%s: %s %s' % (name, state_symbol(worker.get_state()), queue_names)) else: click.echo('worker %s %s %s' % (name, worker.get_state(), queue_names)) else: # Display workers by queue queue_dict = {} for queue in queues: queue_dict[queue] = worker_class.all(queue=queue) if queue_dict: max_length = max([len(q.name) for q, in queue_dict.keys()]) else: max_length = 0 for queue in queue_dict: if queue_dict[queue]: queues_str = ", ".join( sorted( map(lambda w: '%s (%s)' % (w.name, state_symbol(w.get_state())), queue_dict[queue]) ) ) else: queues_str = '–' click.echo('%s %s' % (pad(queue.name + ':', max_length + 1), queues_str)) if not raw: click.echo('%d workers, %d queues' % (len(workers), len(queues))) def show_both(queues, raw, by_queue, queue_class, worker_class): show_queues(queues, raw, by_queue, queue_class, worker_class) if not raw: click.echo('') show_workers(queues, raw, by_queue, queue_class, worker_class) if not raw: click.echo('') import datetime click.echo('Updated: %s' % datetime.datetime.now()) def refresh(interval, func, *args): while True: if interval: click.clear() func(*args) if interval: time.sleep(interval) else: break def setup_loghandlers_from_args(verbose, quiet, date_format, log_format): if verbose and quiet: raise RuntimeError("Flags --verbose and --quiet are mutually exclusive.") if verbose: level = 'DEBUG' elif quiet: level = 'WARNING' else: level = 'INFO' setup_loghandlers(level, date_format=date_format, log_format=log_format) class CliConfig(object): """A helper class to be used with click commands, to handle shared options""" def __init__(self, url=None, config=None, worker_class=DEFAULT_WORKER_CLASS, job_class=DEFAULT_JOB_CLASS, queue_class=DEFAULT_QUEUE_CLASS, connection_class=DEFAULT_CONNECTION_CLASS, path=None, *args, **kwargs): self._connection = None self.url = url self.config = config if path: for pth in path: sys.path.append(pth) try: self.worker_class = import_attribute(worker_class) except (ImportError, AttributeError) as exc: raise click.BadParameter(str(exc), param_hint='--worker-class') try: self.job_class = import_attribute(job_class) except (ImportError, AttributeError) as exc: raise click.BadParameter(str(exc), param_hint='--job-class') try: self.queue_class = import_attribute(queue_class) except (ImportError, AttributeError) as exc: raise click.BadParameter(str(exc), param_hint='--queue-class') try: self.connection_class = import_attribute(connection_class) except (ImportError, AttributeError) as exc: raise click.BadParameter(str(exc), param_hint='--connection-class') @property def connection(self): if self._connection is None: if self.url: self._connection = self.connection_class.from_url(self.url) else: settings = read_config_file(self.config) if self.config else {} self._connection = get_redis_from_config(settings, self.connection_class) return self._connection