Development of an internal social media platform with personalised dashboards for students
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.

amqp.py 11KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380
  1. # -*- coding: utf-8 -*-
  2. """
  3. The :program:`celery amqp` command.
  4. .. program:: celery amqp
  5. """
  6. from __future__ import absolute_import, print_function, unicode_literals
  7. import cmd
  8. import sys
  9. import shlex
  10. import pprint
  11. from functools import partial
  12. from itertools import count
  13. from kombu.utils.encoding import safe_str
  14. from celery.utils.functional import padlist
  15. from celery.bin.base import Command
  16. from celery.five import string_t
  17. from celery.utils import strtobool
  18. __all__ = ['AMQPAdmin', 'AMQShell', 'Spec', 'amqp']
  19. # Map to coerce strings to other types.
  20. COERCE = {bool: strtobool}
  21. HELP_HEADER = """
  22. Commands
  23. --------
  24. """.rstrip()
  25. EXAMPLE_TEXT = """
  26. Example:
  27. -> queue.delete myqueue yes no
  28. """
  29. say = partial(print, file=sys.stderr)
  30. class Spec(object):
  31. """AMQP Command specification.
  32. Used to convert arguments to Python values and display various help
  33. and tooltips.
  34. :param args: see :attr:`args`.
  35. :keyword returns: see :attr:`returns`.
  36. .. attribute args::
  37. List of arguments this command takes. Should
  38. contain `(argument_name, argument_type)` tuples.
  39. .. attribute returns:
  40. Helpful human string representation of what this command returns.
  41. May be :const:`None`, to signify the return type is unknown.
  42. """
  43. def __init__(self, *args, **kwargs):
  44. self.args = args
  45. self.returns = kwargs.get('returns')
  46. def coerce(self, index, value):
  47. """Coerce value for argument at index."""
  48. arg_info = self.args[index]
  49. arg_type = arg_info[1]
  50. # Might be a custom way to coerce the string value,
  51. # so look in the coercion map.
  52. return COERCE.get(arg_type, arg_type)(value)
  53. def str_args_to_python(self, arglist):
  54. """Process list of string arguments to values according to spec.
  55. e.g:
  56. >>> spec = Spec([('queue', str), ('if_unused', bool)])
  57. >>> spec.str_args_to_python('pobox', 'true')
  58. ('pobox', True)
  59. """
  60. return tuple(
  61. self.coerce(index, value) for index, value in enumerate(arglist))
  62. def format_response(self, response):
  63. """Format the return value of this command in a human-friendly way."""
  64. if not self.returns:
  65. return 'ok.' if response is None else response
  66. if callable(self.returns):
  67. return self.returns(response)
  68. return self.returns.format(response)
  69. def format_arg(self, name, type, default_value=None):
  70. if default_value is not None:
  71. return '{0}:{1}'.format(name, default_value)
  72. return name
  73. def format_signature(self):
  74. return ' '.join(self.format_arg(*padlist(list(arg), 3))
  75. for arg in self.args)
  76. def dump_message(message):
  77. if message is None:
  78. return 'No messages in queue. basic.publish something.'
  79. return {'body': message.body,
  80. 'properties': message.properties,
  81. 'delivery_info': message.delivery_info}
  82. def format_declare_queue(ret):
  83. return 'ok. queue:{0} messages:{1} consumers:{2}.'.format(*ret)
  84. class AMQShell(cmd.Cmd):
  85. """AMQP API Shell.
  86. :keyword connect: Function used to connect to the server, must return
  87. connection object.
  88. :keyword silent: If :const:`True`, the commands won't have annoying
  89. output not relevant when running in non-shell mode.
  90. .. attribute: builtins
  91. Mapping of built-in command names -> method names
  92. .. attribute:: amqp
  93. Mapping of AMQP API commands and their :class:`Spec`.
  94. """
  95. conn = None
  96. chan = None
  97. prompt_fmt = '{self.counter}> '
  98. identchars = cmd.IDENTCHARS = '.'
  99. needs_reconnect = False
  100. counter = 1
  101. inc_counter = count(2)
  102. builtins = {'EOF': 'do_exit',
  103. 'exit': 'do_exit',
  104. 'help': 'do_help'}
  105. amqp = {
  106. 'exchange.declare': Spec(('exchange', str),
  107. ('type', str),
  108. ('passive', bool, 'no'),
  109. ('durable', bool, 'no'),
  110. ('auto_delete', bool, 'no'),
  111. ('internal', bool, 'no')),
  112. 'exchange.delete': Spec(('exchange', str),
  113. ('if_unused', bool)),
  114. 'queue.bind': Spec(('queue', str),
  115. ('exchange', str),
  116. ('routing_key', str)),
  117. 'queue.declare': Spec(('queue', str),
  118. ('passive', bool, 'no'),
  119. ('durable', bool, 'no'),
  120. ('exclusive', bool, 'no'),
  121. ('auto_delete', bool, 'no'),
  122. returns=format_declare_queue),
  123. 'queue.delete': Spec(('queue', str),
  124. ('if_unused', bool, 'no'),
  125. ('if_empty', bool, 'no'),
  126. returns='ok. {0} messages deleted.'),
  127. 'queue.purge': Spec(('queue', str),
  128. returns='ok. {0} messages deleted.'),
  129. 'basic.get': Spec(('queue', str),
  130. ('no_ack', bool, 'off'),
  131. returns=dump_message),
  132. 'basic.publish': Spec(('msg', str),
  133. ('exchange', str),
  134. ('routing_key', str),
  135. ('mandatory', bool, 'no'),
  136. ('immediate', bool, 'no')),
  137. 'basic.ack': Spec(('delivery_tag', int)),
  138. }
  139. def _prepare_spec(self, conn):
  140. # XXX Hack to fix Issue #2013
  141. from amqp import Connection, Message
  142. if isinstance(conn.connection, Connection):
  143. self.amqp['basic.publish'] = Spec(('msg', Message),
  144. ('exchange', str),
  145. ('routing_key', str),
  146. ('mandatory', bool, 'no'),
  147. ('immediate', bool, 'no'))
  148. def __init__(self, *args, **kwargs):
  149. self.connect = kwargs.pop('connect')
  150. self.silent = kwargs.pop('silent', False)
  151. self.out = kwargs.pop('out', sys.stderr)
  152. cmd.Cmd.__init__(self, *args, **kwargs)
  153. self._reconnect()
  154. def note(self, m):
  155. """Say something to the user. Disabled if :attr:`silent`."""
  156. if not self.silent:
  157. say(m, file=self.out)
  158. def say(self, m):
  159. say(m, file=self.out)
  160. def get_amqp_api_command(self, cmd, arglist):
  161. """With a command name and a list of arguments, convert the arguments
  162. to Python values and find the corresponding method on the AMQP channel
  163. object.
  164. :returns: tuple of `(method, processed_args)`.
  165. """
  166. spec = self.amqp[cmd]
  167. args = spec.str_args_to_python(arglist)
  168. attr_name = cmd.replace('.', '_')
  169. if self.needs_reconnect:
  170. self._reconnect()
  171. return getattr(self.chan, attr_name), args, spec.format_response
  172. def do_exit(self, *args):
  173. """The `'exit'` command."""
  174. self.note("\n-> please, don't leave!")
  175. sys.exit(0)
  176. def display_command_help(self, cmd, short=False):
  177. spec = self.amqp[cmd]
  178. self.say('{0} {1}'.format(cmd, spec.format_signature()))
  179. def do_help(self, *args):
  180. if not args:
  181. self.say(HELP_HEADER)
  182. for cmd_name in self.amqp:
  183. self.display_command_help(cmd_name, short=True)
  184. self.say(EXAMPLE_TEXT)
  185. else:
  186. self.display_command_help(args[0])
  187. def default(self, line):
  188. self.say("unknown syntax: {0!r}. how about some 'help'?".format(line))
  189. def get_names(self):
  190. return set(self.builtins) | set(self.amqp)
  191. def completenames(self, text, *ignored):
  192. """Return all commands starting with `text`, for tab-completion."""
  193. names = self.get_names()
  194. first = [cmd for cmd in names
  195. if cmd.startswith(text.replace('_', '.'))]
  196. if first:
  197. return first
  198. return [cmd for cmd in names
  199. if cmd.partition('.')[2].startswith(text)]
  200. def dispatch(self, cmd, argline):
  201. """Dispatch and execute the command.
  202. Lookup order is: :attr:`builtins` -> :attr:`amqp`.
  203. """
  204. arglist = shlex.split(safe_str(argline))
  205. if cmd in self.builtins:
  206. return getattr(self, self.builtins[cmd])(*arglist)
  207. fun, args, formatter = self.get_amqp_api_command(cmd, arglist)
  208. return formatter(fun(*args))
  209. def parseline(self, line):
  210. """Parse input line.
  211. :returns: tuple of three items:
  212. `(command_name, arglist, original_line)`
  213. """
  214. parts = line.split()
  215. if parts:
  216. return parts[0], ' '.join(parts[1:]), line
  217. return '', '', line
  218. def onecmd(self, line):
  219. """Parse line and execute command."""
  220. cmd, arg, line = self.parseline(line)
  221. if not line:
  222. return self.emptyline()
  223. self.lastcmd = line
  224. self.counter = next(self.inc_counter)
  225. try:
  226. self.respond(self.dispatch(cmd, arg))
  227. except (AttributeError, KeyError) as exc:
  228. self.default(line)
  229. except Exception as exc:
  230. self.say(exc)
  231. self.needs_reconnect = True
  232. def respond(self, retval):
  233. """What to do with the return value of a command."""
  234. if retval is not None:
  235. if isinstance(retval, string_t):
  236. self.say(retval)
  237. else:
  238. self.say(pprint.pformat(retval))
  239. def _reconnect(self):
  240. """Re-establish connection to the AMQP server."""
  241. self.conn = self.connect(self.conn)
  242. self._prepare_spec(self.conn)
  243. self.chan = self.conn.default_channel
  244. self.needs_reconnect = False
  245. @property
  246. def prompt(self):
  247. return self.prompt_fmt.format(self=self)
  248. class AMQPAdmin(object):
  249. """The celery :program:`celery amqp` utility."""
  250. Shell = AMQShell
  251. def __init__(self, *args, **kwargs):
  252. self.app = kwargs['app']
  253. self.out = kwargs.setdefault('out', sys.stderr)
  254. self.silent = kwargs.get('silent')
  255. self.args = args
  256. def connect(self, conn=None):
  257. if conn:
  258. conn.close()
  259. conn = self.app.connection()
  260. self.note('-> connecting to {0}.'.format(conn.as_uri()))
  261. conn.connect()
  262. self.note('-> connected.')
  263. return conn
  264. def run(self):
  265. shell = self.Shell(connect=self.connect, out=self.out)
  266. if self.args:
  267. return shell.onecmd(' '.join(self.args))
  268. try:
  269. return shell.cmdloop()
  270. except KeyboardInterrupt:
  271. self.note('(bibi)')
  272. pass
  273. def note(self, m):
  274. if not self.silent:
  275. say(m, file=self.out)
  276. class amqp(Command):
  277. """AMQP Administration Shell.
  278. Also works for non-amqp transports (but not ones that
  279. store declarations in memory).
  280. Examples::
  281. celery amqp
  282. start shell mode
  283. celery amqp help
  284. show list of commands
  285. celery amqp exchange.delete name
  286. celery amqp queue.delete queue
  287. celery amqp queue.delete queue yes yes
  288. """
  289. def run(self, *args, **options):
  290. options['app'] = self.app
  291. return AMQPAdmin(*args, **options).run()
  292. def main():
  293. amqp().execute_from_commandline()
  294. if __name__ == '__main__': # pragma: no cover
  295. main()