You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

238 lines
7.7 KiB

  1. # -*- coding: utf-8 -*-
  2. from __future__ import (absolute_import, division, print_function,
  3. unicode_literals)
  4. import sys
  5. import importlib
  6. import time
  7. from functools import partial
  8. import click
  9. import redis
  10. from redis import Redis
  11. from redis.sentinel import Sentinel
  12. from rq.defaults import (DEFAULT_CONNECTION_CLASS, DEFAULT_JOB_CLASS,
  13. DEFAULT_QUEUE_CLASS, DEFAULT_WORKER_CLASS)
  14. from rq.logutils import setup_loghandlers
  15. from rq.utils import import_attribute
  16. from rq.worker import WorkerStatus
  17. red = partial(click.style, fg='red')
  18. green = partial(click.style, fg='green')
  19. yellow = partial(click.style, fg='yellow')
  20. def read_config_file(module):
  21. """Reads all UPPERCASE variables defined in the given module file."""
  22. settings = importlib.import_module(module)
  23. return dict([(k, v)
  24. for k, v in settings.__dict__.items()
  25. if k.upper() == k])
  26. def get_redis_from_config(settings, connection_class=Redis):
  27. """Returns a StrictRedis instance from a dictionary of settings.
  28. To use redis sentinel, you must specify a dictionary in the configuration file.
  29. Example of a dictionary with keys without values:
  30. SENTINEL: {'INSTANCES':, 'SOCKET_TIMEOUT':, 'PASSWORD':,'DB':, 'MASTER_NAME':}
  31. """
  32. if settings.get('REDIS_URL') is not None:
  33. return connection_class.from_url(settings['REDIS_URL'])
  34. elif settings.get('SENTINEL') is not None:
  35. instances = settings['SENTINEL'].get('INSTANCES', [('localhost', 26379)])
  36. socket_timeout = settings['SENTINEL'].get('SOCKET_TIMEOUT', None)
  37. password = settings['SENTINEL'].get('PASSWORD', None)
  38. db = settings['SENTINEL'].get('DB', 0)
  39. master_name = settings['SENTINEL'].get('MASTER_NAME', 'mymaster')
  40. sn = Sentinel(instances, socket_timeout=socket_timeout, password=password, db=db)
  41. return sn.master_for(master_name)
  42. kwargs = {
  43. 'host': settings.get('REDIS_HOST', 'localhost'),
  44. 'port': settings.get('REDIS_PORT', 6379),
  45. 'db': settings.get('REDIS_DB', 0),
  46. 'password': settings.get('REDIS_PASSWORD', None),
  47. 'ssl': settings.get('REDIS_SSL', False),
  48. }
  49. return connection_class(**kwargs)
  50. def pad(s, pad_to_length):
  51. """Pads the given string to the given length."""
  52. return ('%-' + '%ds' % pad_to_length) % (s,)
  53. def get_scale(x):
  54. """Finds the lowest scale where x <= scale."""
  55. scales = [20, 50, 100, 200, 400, 600, 800, 1000]
  56. for scale in scales:
  57. if x <= scale:
  58. return scale
  59. return x
  60. def state_symbol(state):
  61. symbols = {
  62. WorkerStatus.BUSY: red('busy'),
  63. WorkerStatus.IDLE: green('idle'),
  64. WorkerStatus.SUSPENDED: yellow('suspended'),
  65. }
  66. try:
  67. return symbols[state]
  68. except KeyError:
  69. return state
  70. def show_queues(queues, raw, by_queue, queue_class, worker_class):
  71. num_jobs = 0
  72. termwidth, _ = click.get_terminal_size()
  73. chartwidth = min(20, termwidth - 20)
  74. max_count = 0
  75. counts = dict()
  76. for q in queues:
  77. count = q.count
  78. counts[q] = count
  79. max_count = max(max_count, count)
  80. scale = get_scale(max_count)
  81. ratio = chartwidth * 1.0 / scale
  82. for q in queues:
  83. count = counts[q]
  84. if not raw:
  85. chart = green('|' + '█' * int(ratio * count))
  86. line = '%-12s %s %d' % (q.name, chart, count)
  87. else:
  88. line = 'queue %s %d' % (q.name, count)
  89. click.echo(line)
  90. num_jobs += count
  91. # print summary when not in raw mode
  92. if not raw:
  93. click.echo('%d queues, %d jobs total' % (len(queues), num_jobs))
  94. def show_workers(queues, raw, by_queue, queue_class, worker_class):
  95. workers = set()
  96. for queue in queues:
  97. for worker in worker_class.all(queue=queue):
  98. workers.add(worker)
  99. if not by_queue:
  100. for worker in workers:
  101. queue_names = ', '.join(worker.queue_names())
  102. name = '%s (%s %s)' % (worker.name, worker.hostname, worker.pid)
  103. if not raw:
  104. click.echo('%s: %s %s' % (name, state_symbol(worker.get_state()), queue_names))
  105. else:
  106. click.echo('worker %s %s %s' % (name, worker.get_state(), queue_names))
  107. else:
  108. # Display workers by queue
  109. queue_dict = {}
  110. for queue in queues:
  111. queue_dict[queue] = worker_class.all(queue=queue)
  112. if queue_dict:
  113. max_length = max([len(q.name) for q, in queue_dict.keys()])
  114. else:
  115. max_length = 0
  116. for queue in queue_dict:
  117. if queue_dict[queue]:
  118. queues_str = ", ".join(
  119. sorted(
  120. map(lambda w: '%s (%s)' % (w.name, state_symbol(w.get_state())), queue_dict[queue])
  121. )
  122. )
  123. else:
  124. queues_str = '–'
  125. click.echo('%s %s' % (pad(queue.name + ':', max_length + 1), queues_str))
  126. if not raw:
  127. click.echo('%d workers, %d queues' % (len(workers), len(queues)))
  128. def show_both(queues, raw, by_queue, queue_class, worker_class):
  129. show_queues(queues, raw, by_queue, queue_class, worker_class)
  130. if not raw:
  131. click.echo('')
  132. show_workers(queues, raw, by_queue, queue_class, worker_class)
  133. if not raw:
  134. click.echo('')
  135. import datetime
  136. click.echo('Updated: %s' % datetime.datetime.now())
  137. def refresh(interval, func, *args):
  138. while True:
  139. if interval:
  140. click.clear()
  141. func(*args)
  142. if interval:
  143. time.sleep(interval)
  144. else:
  145. break
  146. def setup_loghandlers_from_args(verbose, quiet, date_format, log_format):
  147. if verbose and quiet:
  148. raise RuntimeError("Flags --verbose and --quiet are mutually exclusive.")
  149. if verbose:
  150. level = 'DEBUG'
  151. elif quiet:
  152. level = 'WARNING'
  153. else:
  154. level = 'INFO'
  155. setup_loghandlers(level, date_format=date_format, log_format=log_format)
  156. class CliConfig(object):
  157. """A helper class to be used with click commands, to handle shared options"""
  158. def __init__(self, url=None, config=None, worker_class=DEFAULT_WORKER_CLASS,
  159. job_class=DEFAULT_JOB_CLASS, queue_class=DEFAULT_QUEUE_CLASS,
  160. connection_class=DEFAULT_CONNECTION_CLASS, path=None, *args, **kwargs):
  161. self._connection = None
  162. self.url = url
  163. self.config = config
  164. if path:
  165. for pth in path:
  166. sys.path.append(pth)
  167. try:
  168. self.worker_class = import_attribute(worker_class)
  169. except (ImportError, AttributeError) as exc:
  170. raise click.BadParameter(str(exc), param_hint='--worker-class')
  171. try:
  172. self.job_class = import_attribute(job_class)
  173. except (ImportError, AttributeError) as exc:
  174. raise click.BadParameter(str(exc), param_hint='--job-class')
  175. try:
  176. self.queue_class = import_attribute(queue_class)
  177. except (ImportError, AttributeError) as exc:
  178. raise click.BadParameter(str(exc), param_hint='--queue-class')
  179. try:
  180. self.connection_class = import_attribute(connection_class)
  181. except (ImportError, AttributeError) as exc:
  182. raise click.BadParameter(str(exc), param_hint='--connection-class')
  183. @property
  184. def connection(self):
  185. if self._connection is None:
  186. if self.url:
  187. self._connection = self.connection_class.from_url(self.url)
  188. else:
  189. settings = read_config_file(self.config) if self.config else {}
  190. self._connection = get_redis_from_config(settings,
  191. self.connection_class)
  192. return self._connection