base.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611
  1. # -*- coding: utf-8 -*-
  2. """
  3. .. _preload-options:
  4. Preload Options
  5. ---------------
  6. These options are supported by all commands,
  7. and usually parsed before command-specific arguments.
  8. .. cmdoption:: -A, --app
  9. app instance to use (e.g. module.attr_name)
  10. .. cmdoption:: -b, --broker
  11. url to broker. default is 'amqp://guest@localhost//'
  12. .. cmdoption:: --loader
  13. name of custom loader class to use.
  14. .. cmdoption:: --config
  15. Name of the configuration module
  16. .. _daemon-options:
  17. Daemon Options
  18. --------------
  19. These options are supported by commands that can detach
  20. into the background (daemon). They will be present
  21. in any command that also has a `--detach` option.
  22. .. cmdoption:: -f, --logfile
  23. Path to log file. If no logfile is specified, `stderr` is used.
  24. .. cmdoption:: --pidfile
  25. Optional file used to store the process pid.
  26. The program will not start if this file already exists
  27. and the pid is still alive.
  28. .. cmdoption:: --uid
  29. User id, or user name of the user to run as after detaching.
  30. .. cmdoption:: --gid
  31. Group id, or group name of the main group to change to after
  32. detaching.
  33. .. cmdoption:: --umask
  34. Effective umask of the process after detaching. Default is 0.
  35. .. cmdoption:: --workdir
  36. Optional directory to change to after detaching.
  37. """
  38. from __future__ import absolute_import, print_function
  39. import os
  40. import re
  41. import socket
  42. import sys
  43. import warnings
  44. import json
  45. from collections import defaultdict
  46. from heapq import heappush
  47. from inspect import getargspec
  48. from optparse import OptionParser, IndentedHelpFormatter, make_option as Option
  49. from pprint import pformat
  50. from types import ModuleType
  51. from celery import VERSION_BANNER, Celery, maybe_patch_concurrency
  52. from celery import signals
  53. from celery.exceptions import CDeprecationWarning, CPendingDeprecationWarning
  54. from celery.five import items, string, string_t, values
  55. from celery.platforms import EX_FAILURE, EX_OK, EX_USAGE
  56. from celery.utils import term
  57. from celery.utils import text
  58. from celery.utils.imports import symbol_by_name, import_from_cwd
  59. # always enable DeprecationWarnings, so our users can see them.
  60. for warning in (CDeprecationWarning, CPendingDeprecationWarning):
  61. warnings.simplefilter('once', warning, 0)
  62. ARGV_DISABLED = """
  63. Unrecognized command-line arguments: {0}
  64. Try --help?
  65. """
  66. find_long_opt = re.compile(r'.+?(--.+?)(?:\s|,|$)')
  67. find_rst_ref = re.compile(r':\w+:`(.+?)`')
  68. find_sformat = re.compile(r'%(\w)')
  69. __all__ = ['Error', 'UsageError', 'Extensions', 'HelpFormatter',
  70. 'Command', 'Option', 'daemon_options']
  71. class Error(Exception):
  72. status = EX_FAILURE
  73. def __init__(self, reason, status=None):
  74. self.reason = reason
  75. self.status = status if status is not None else self.status
  76. super(Error, self).__init__(reason, status)
  77. def __str__(self):
  78. return self.reason
  79. __unicode__ = __str__
  80. class UsageError(Error):
  81. status = EX_USAGE
  82. class Extensions(object):
  83. def __init__(self, namespace, register):
  84. self.names = []
  85. self.namespace = namespace
  86. self.register = register
  87. def add(self, cls, name):
  88. heappush(self.names, name)
  89. self.register(cls, name=name)
  90. def load(self):
  91. try:
  92. from pkg_resources import iter_entry_points
  93. except ImportError: # pragma: no cover
  94. return
  95. for ep in iter_entry_points(self.namespace):
  96. sym = ':'.join([ep.module_name, ep.attrs[0]])
  97. try:
  98. cls = symbol_by_name(sym)
  99. except (ImportError, SyntaxError) as exc:
  100. warnings.warn(
  101. 'Cannot load extension {0!r}: {1!r}'.format(sym, exc))
  102. else:
  103. self.add(cls, ep.name)
  104. return self.names
  105. class HelpFormatter(IndentedHelpFormatter):
  106. def format_epilog(self, epilog):
  107. if epilog:
  108. return '\n{0}\n\n'.format(epilog)
  109. return ''
  110. def format_description(self, description):
  111. return text.ensure_2lines(text.fill_paragraphs(
  112. text.dedent(description), self.width))
  113. class Command(object):
  114. """Base class for command-line applications.
  115. :keyword app: The current app.
  116. :keyword get_app: Callable returning the current app if no app provided.
  117. """
  118. Error = Error
  119. UsageError = UsageError
  120. Parser = OptionParser
  121. #: Arg list used in help.
  122. args = ''
  123. #: Application version.
  124. version = VERSION_BANNER
  125. #: If false the parser will raise an exception if positional
  126. #: args are provided.
  127. supports_args = True
  128. #: List of options (without preload options).
  129. option_list = ()
  130. # module Rst documentation to parse help from (if any)
  131. doc = None
  132. # Some programs (multi) does not want to load the app specified
  133. # (Issue #1008).
  134. respects_app_option = True
  135. #: List of options to parse before parsing other options.
  136. preload_options = (
  137. Option('-A', '--app', default=None),
  138. Option('-b', '--broker', default=None),
  139. Option('--loader', default=None),
  140. Option('--config', default=None),
  141. Option('--workdir', default=None, dest='working_directory'),
  142. Option('--no-color', '-C', action='store_true', default=None),
  143. Option('--quiet', '-q', action='store_true'),
  144. )
  145. #: Enable if the application should support config from the cmdline.
  146. enable_config_from_cmdline = False
  147. #: Default configuration namespace.
  148. namespace = 'celery'
  149. #: Text to print at end of --help
  150. epilog = None
  151. #: Text to print in --help before option list.
  152. description = ''
  153. #: Set to true if this command doesn't have subcommands
  154. leaf = True
  155. # used by :meth:`say_remote_command_reply`.
  156. show_body = True
  157. # used by :meth:`say_chat`.
  158. show_reply = True
  159. prog_name = 'celery'
  160. def __init__(self, app=None, get_app=None, no_color=False,
  161. stdout=None, stderr=None, quiet=False, on_error=None,
  162. on_usage_error=None):
  163. self.app = app
  164. self.get_app = get_app or self._get_default_app
  165. self.stdout = stdout or sys.stdout
  166. self.stderr = stderr or sys.stderr
  167. self.no_color = no_color
  168. self.colored = term.colored(enabled=not self.no_color)
  169. self.quiet = quiet
  170. if not self.description:
  171. self.description = self.__doc__
  172. if on_error:
  173. self.on_error = on_error
  174. if on_usage_error:
  175. self.on_usage_error = on_usage_error
  176. def run(self, *args, **options):
  177. """This is the body of the command called by :meth:`handle_argv`."""
  178. raise NotImplementedError('subclass responsibility')
  179. def on_error(self, exc):
  180. self.error(self.colored.red('Error: {0}'.format(exc)))
  181. def on_usage_error(self, exc):
  182. self.handle_error(exc)
  183. def on_concurrency_setup(self):
  184. pass
  185. def __call__(self, *args, **kwargs):
  186. self.verify_args(args)
  187. try:
  188. ret = self.run(*args, **kwargs)
  189. return ret if ret is not None else EX_OK
  190. except self.UsageError as exc:
  191. self.on_usage_error(exc)
  192. return exc.status
  193. except self.Error as exc:
  194. self.on_error(exc)
  195. return exc.status
  196. def verify_args(self, given, _index=0):
  197. S = getargspec(self.run)
  198. _index = 1 if S.args and S.args[0] == 'self' else _index
  199. required = S.args[_index:-len(S.defaults) if S.defaults else None]
  200. missing = required[len(given):]
  201. if missing:
  202. raise self.UsageError('Missing required {0}: {1}'.format(
  203. text.pluralize(len(missing), 'argument'),
  204. ', '.join(missing)
  205. ))
  206. def execute_from_commandline(self, argv=None):
  207. """Execute application from command-line.
  208. :keyword argv: The list of command-line arguments.
  209. Defaults to ``sys.argv``.
  210. """
  211. if argv is None:
  212. argv = list(sys.argv)
  213. # Should we load any special concurrency environment?
  214. self.maybe_patch_concurrency(argv)
  215. self.on_concurrency_setup()
  216. # Dump version and exit if '--version' arg set.
  217. self.early_version(argv)
  218. argv = self.setup_app_from_commandline(argv)
  219. self.prog_name = os.path.basename(argv[0])
  220. return self.handle_argv(self.prog_name, argv[1:])
  221. def run_from_argv(self, prog_name, argv=None, command=None):
  222. return self.handle_argv(prog_name,
  223. sys.argv if argv is None else argv, command)
  224. def maybe_patch_concurrency(self, argv=None):
  225. argv = argv or sys.argv
  226. pool_option = self.with_pool_option(argv)
  227. if pool_option:
  228. maybe_patch_concurrency(argv, *pool_option)
  229. short_opts, long_opts = pool_option
  230. def usage(self, command):
  231. return '%prog {0} [options] {self.args}'.format(command, self=self)
  232. def get_options(self):
  233. """Get supported command-line options."""
  234. return self.option_list
  235. def expanduser(self, value):
  236. if isinstance(value, string_t):
  237. return os.path.expanduser(value)
  238. return value
  239. def handle_argv(self, prog_name, argv, command=None):
  240. """Parse command-line arguments from ``argv`` and dispatch
  241. to :meth:`run`.
  242. :param prog_name: The program name (``argv[0]``).
  243. :param argv: Command arguments.
  244. Exits with an error message if :attr:`supports_args` is disabled
  245. and ``argv`` contains positional arguments.
  246. """
  247. options, args = self.prepare_args(
  248. *self.parse_options(prog_name, argv, command))
  249. return self(*args, **options)
  250. def prepare_args(self, options, args):
  251. if options:
  252. options = dict((k, self.expanduser(v))
  253. for k, v in items(vars(options))
  254. if not k.startswith('_'))
  255. args = [self.expanduser(arg) for arg in args]
  256. self.check_args(args)
  257. return options, args
  258. def check_args(self, args):
  259. if not self.supports_args and args:
  260. self.die(ARGV_DISABLED.format(', '.join(args)), EX_USAGE)
  261. def error(self, s):
  262. self.out(s, fh=self.stderr)
  263. def out(self, s, fh=None):
  264. print(s, file=fh or self.stdout)
  265. def die(self, msg, status=EX_FAILURE):
  266. self.error(msg)
  267. sys.exit(status)
  268. def early_version(self, argv):
  269. if '--version' in argv:
  270. print(self.version, file=self.stdout)
  271. sys.exit(0)
  272. def parse_options(self, prog_name, arguments, command=None):
  273. """Parse the available options."""
  274. # Don't want to load configuration to just print the version,
  275. # so we handle --version manually here.
  276. self.parser = self.create_parser(prog_name, command)
  277. return self.parser.parse_args(arguments)
  278. def create_parser(self, prog_name, command=None):
  279. option_list = (
  280. self.preload_options +
  281. self.get_options() +
  282. tuple(self.app.user_options['preload'])
  283. )
  284. return self.prepare_parser(self.Parser(
  285. prog=prog_name,
  286. usage=self.usage(command),
  287. version=self.version,
  288. epilog=self.epilog,
  289. formatter=HelpFormatter(),
  290. description=self.description,
  291. option_list=option_list,
  292. ))
  293. def prepare_parser(self, parser):
  294. docs = [self.parse_doc(doc) for doc in (self.doc, __doc__) if doc]
  295. for doc in docs:
  296. for long_opt, help in items(doc):
  297. option = parser.get_option(long_opt)
  298. if option is not None:
  299. option.help = ' '.join(help).format(default=option.default)
  300. return parser
  301. def setup_app_from_commandline(self, argv):
  302. preload_options = self.parse_preload_options(argv)
  303. quiet = preload_options.get('quiet')
  304. if quiet is not None:
  305. self.quiet = quiet
  306. self.colored.enabled = \
  307. not preload_options.get('no_color', self.no_color)
  308. workdir = preload_options.get('working_directory')
  309. if workdir:
  310. os.chdir(workdir)
  311. app = (preload_options.get('app') or
  312. os.environ.get('CELERY_APP') or
  313. self.app)
  314. preload_loader = preload_options.get('loader')
  315. if preload_loader:
  316. # Default app takes loader from this env (Issue #1066).
  317. os.environ['CELERY_LOADER'] = preload_loader
  318. loader = (preload_loader,
  319. os.environ.get('CELERY_LOADER') or
  320. 'default')
  321. broker = preload_options.get('broker', None)
  322. if broker:
  323. os.environ['CELERY_BROKER_URL'] = broker
  324. config = preload_options.get('config')
  325. if config:
  326. os.environ['CELERY_CONFIG_MODULE'] = config
  327. if self.respects_app_option:
  328. if app and self.respects_app_option:
  329. self.app = self.find_app(app)
  330. elif self.app is None:
  331. self.app = self.get_app(loader=loader)
  332. if self.enable_config_from_cmdline:
  333. argv = self.process_cmdline_config(argv)
  334. else:
  335. self.app = Celery()
  336. user_preload = tuple(self.app.user_options['preload'] or ())
  337. if user_preload:
  338. user_options = self.preparse_options(argv, user_preload)
  339. for user_option in user_preload:
  340. user_options.setdefault(user_option.dest, user_option.default)
  341. signals.user_preload_options.send(
  342. sender=self, app=self.app, options=user_options,
  343. )
  344. return argv
  345. def find_app(self, app):
  346. try:
  347. sym = self.symbol_by_name(app)
  348. except AttributeError:
  349. # last part was not an attribute, but a module
  350. sym = import_from_cwd(app)
  351. if isinstance(sym, ModuleType):
  352. try:
  353. return sym.app
  354. except AttributeError:
  355. try:
  356. return sym.celery
  357. except AttributeError:
  358. if getattr(sym, '__path__', None):
  359. try:
  360. return self.find_app(
  361. '{0}.celery:'.format(app.replace(':', '')),
  362. )
  363. except ImportError:
  364. pass
  365. for suspect in values(vars(sym)):
  366. if isinstance(suspect, Celery):
  367. return suspect
  368. raise
  369. return sym
  370. def symbol_by_name(self, name):
  371. return symbol_by_name(name, imp=import_from_cwd)
  372. get_cls_by_name = symbol_by_name # XXX compat
  373. def process_cmdline_config(self, argv):
  374. try:
  375. cargs_start = argv.index('--')
  376. except ValueError:
  377. return argv
  378. argv, cargs = argv[:cargs_start], argv[cargs_start + 1:]
  379. self.app.config_from_cmdline(cargs, namespace=self.namespace)
  380. return argv
  381. def parse_preload_options(self, args):
  382. return self.preparse_options(args, self.preload_options)
  383. def preparse_options(self, args, options):
  384. acc = {}
  385. opts = {}
  386. for opt in options:
  387. for t in (opt._long_opts, opt._short_opts):
  388. opts.update(dict(zip(t, [opt] * len(t))))
  389. index = 0
  390. length = len(args)
  391. while index < length:
  392. arg = args[index]
  393. if arg.startswith('--'):
  394. if '=' in arg:
  395. key, value = arg.split('=', 1)
  396. opt = opts.get(key)
  397. if opt:
  398. acc[opt.dest] = value
  399. else:
  400. opt = opts.get(arg)
  401. if opt and opt.action == 'store_true':
  402. acc[opt.dest] = True
  403. elif arg.startswith('-'):
  404. opt = opts.get(arg)
  405. if opt:
  406. if opt.takes_value():
  407. acc[opt.dest] = args[index + 1]
  408. index += 1
  409. elif opt.action == 'store_true':
  410. acc[opt.dest] = True
  411. index += 1
  412. return acc
  413. def parse_doc(self, doc):
  414. options, in_option = defaultdict(list), None
  415. for line in doc.splitlines():
  416. if line.startswith('.. cmdoption::'):
  417. m = find_long_opt.match(line)
  418. if m:
  419. in_option = m.groups()[0].strip()
  420. assert in_option, 'missing long opt'
  421. elif in_option and line.startswith(' ' * 4):
  422. options[in_option].append(
  423. find_rst_ref.sub(r'\1', line.strip()).replace('`', ''))
  424. return options
  425. def with_pool_option(self, argv):
  426. """Return tuple of ``(short_opts, long_opts)`` if the command
  427. supports a pool argument, and used to monkey patch eventlet/gevent
  428. environments as early as possible.
  429. E.g::
  430. has_pool_option = (['-P'], ['--pool'])
  431. """
  432. pass
  433. def simple_format(self, s, match=find_sformat, expand=r'\1', **keys):
  434. if s:
  435. host = socket.gethostname()
  436. name, _, domain = host.partition('.')
  437. keys = dict({'%': '%', 'h': host, 'n': name, 'd': domain}, **keys)
  438. return match.sub(lambda m: keys[m.expand(expand)], s)
  439. return s
  440. def _get_default_app(self, *args, **kwargs):
  441. from celery._state import get_current_app
  442. return get_current_app() # omit proxy
  443. def pretty_list(self, n):
  444. c = self.colored
  445. if not n:
  446. return '- empty -'
  447. return '\n'.join(
  448. str(c.reset(c.white('*'), ' {0}'.format(item))) for item in n
  449. )
  450. def pretty_dict_ok_error(self, n):
  451. c = self.colored
  452. try:
  453. return (c.green('OK'),
  454. text.indent(self.pretty(n['ok'])[1], 4))
  455. except KeyError:
  456. pass
  457. return (c.red('ERROR'),
  458. text.indent(self.pretty(n['error'])[1], 4))
  459. def say_remote_command_reply(self, replies):
  460. c = self.colored
  461. node = next(iter(replies)) # <-- take first.
  462. reply = replies[node]
  463. status, preply = self.pretty(reply)
  464. self.say_chat('->', c.cyan(node, ': ') + status,
  465. text.indent(preply, 4) if self.show_reply else '')
  466. def pretty(self, n):
  467. OK = str(self.colored.green('OK'))
  468. if isinstance(n, list):
  469. return OK, self.pretty_list(n)
  470. if isinstance(n, dict):
  471. if 'ok' in n or 'error' in n:
  472. return self.pretty_dict_ok_error(n)
  473. else:
  474. return OK, json.dumps(n, sort_keys=True, indent=4)
  475. if isinstance(n, string_t):
  476. return OK, string(n)
  477. return OK, pformat(n)
  478. def say_chat(self, direction, title, body=''):
  479. c = self.colored
  480. if direction == '<-' and self.quiet:
  481. return
  482. dirstr = not self.quiet and c.bold(c.white(direction), ' ') or ''
  483. self.out(c.reset(dirstr, title))
  484. if body and self.show_body:
  485. self.out(body)
  486. def daemon_options(default_pidfile=None, default_logfile=None):
  487. return (
  488. Option('-f', '--logfile', default=default_logfile),
  489. Option('--pidfile', default=default_pidfile),
  490. Option('--uid', default=None),
  491. Option('--gid', default=None),
  492. Option('--umask', default=0, type='int'),
  493. )