amqp.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631
  1. # -*- coding: utf-8 -*-
  2. """Sending/Receiving Messages (Kombu integration)."""
  3. from __future__ import absolute_import, unicode_literals
  4. import numbers
  5. import sys
  6. from collections import Mapping, namedtuple
  7. from datetime import timedelta
  8. from weakref import WeakValueDictionary
  9. from celery import signals
  10. from celery.five import items, string_t
  11. from celery.local import try_import
  12. from celery.utils.nodenames import anon_nodename
  13. from celery.utils.saferepr import saferepr
  14. from celery.utils.text import indent as textindent
  15. from celery.utils.time import maybe_make_aware
  16. from kombu import Connection, Consumer, Exchange, Producer, Queue, pools
  17. from kombu.common import Broadcast
  18. from kombu.utils.functional import maybe_list
  19. from kombu.utils.objects import cached_property
  20. from . import routes as _routes
  21. __all__ = ('AMQP', 'Queues', 'task_message')
  22. PY3 = sys.version_info[0] == 3
  23. #: earliest date supported by time.mktime.
  24. INT_MIN = -2147483648
  25. # json in Python 2.7 borks if dict contains byte keys.
  26. JSON_NEEDS_UNICODE_KEYS = not PY3 and not try_import('simplejson')
  27. #: Human readable queue declaration.
  28. QUEUE_FORMAT = """
  29. .> {0.name:<16} exchange={0.exchange.name}({0.exchange.type}) \
  30. key={0.routing_key}
  31. """
  32. task_message = namedtuple('task_message',
  33. ('headers', 'properties', 'body', 'sent_event'))
  34. def utf8dict(d, encoding='utf-8'):
  35. return {k.decode(encoding) if isinstance(k, bytes) else k: v
  36. for k, v in items(d)}
  37. class Queues(dict):
  38. """Queue name⇒ declaration mapping.
  39. Arguments:
  40. queues (Iterable): Initial list/tuple or dict of queues.
  41. create_missing (bool): By default any unknown queues will be
  42. added automatically, but if this flag is disabled the occurrence
  43. of unknown queues in `wanted` will raise :exc:`KeyError`.
  44. ha_policy (Sequence, str): Default HA policy for queues with none set.
  45. max_priority (int): Default x-max-priority for queues with none set.
  46. """
  47. #: If set, this is a subset of queues to consume from.
  48. #: The rest of the queues are then used for routing only.
  49. _consume_from = None
  50. def __init__(self, queues=None, default_exchange=None,
  51. create_missing=True, ha_policy=None, autoexchange=None,
  52. max_priority=None, default_routing_key=None):
  53. dict.__init__(self)
  54. self.aliases = WeakValueDictionary()
  55. self.default_exchange = default_exchange
  56. self.default_routing_key = default_routing_key
  57. self.create_missing = create_missing
  58. self.ha_policy = ha_policy
  59. self.autoexchange = Exchange if autoexchange is None else autoexchange
  60. self.max_priority = max_priority
  61. if queues is not None and not isinstance(queues, Mapping):
  62. queues = {q.name: q for q in queues}
  63. for name, q in items(queues or {}):
  64. self.add(q) if isinstance(q, Queue) else self.add_compat(name, **q)
  65. def __getitem__(self, name):
  66. try:
  67. return self.aliases[name]
  68. except KeyError:
  69. return dict.__getitem__(self, name)
  70. def __setitem__(self, name, queue):
  71. if self.default_exchange and not queue.exchange:
  72. queue.exchange = self.default_exchange
  73. dict.__setitem__(self, name, queue)
  74. if queue.alias:
  75. self.aliases[queue.alias] = queue
  76. def __missing__(self, name):
  77. if self.create_missing:
  78. return self.add(self.new_missing(name))
  79. raise KeyError(name)
  80. def add(self, queue, **kwargs):
  81. """Add new queue.
  82. The first argument can either be a :class:`kombu.Queue` instance,
  83. or the name of a queue. If the former the rest of the keyword
  84. arguments are ignored, and options are simply taken from the queue
  85. instance.
  86. Arguments:
  87. queue (kombu.Queue, str): Queue to add.
  88. exchange (kombu.Exchange, str):
  89. if queue is str, specifies exchange name.
  90. routing_key (str): if queue is str, specifies binding key.
  91. exchange_type (str): if queue is str, specifies type of exchange.
  92. **options (Any): Additional declaration options used when
  93. queue is a str.
  94. """
  95. if not isinstance(queue, Queue):
  96. return self.add_compat(queue, **kwargs)
  97. return self._add(queue)
  98. def add_compat(self, name, **options):
  99. # docs used to use binding_key as routing key
  100. options.setdefault('routing_key', options.get('binding_key'))
  101. if options['routing_key'] is None:
  102. options['routing_key'] = name
  103. return self._add(Queue.from_dict(name, **options))
  104. def _add(self, queue):
  105. if not queue.routing_key:
  106. if queue.exchange is None or queue.exchange.name == '':
  107. queue.exchange = self.default_exchange
  108. queue.routing_key = self.default_routing_key
  109. if self.ha_policy:
  110. if queue.queue_arguments is None:
  111. queue.queue_arguments = {}
  112. self._set_ha_policy(queue.queue_arguments)
  113. if self.max_priority is not None:
  114. if queue.queue_arguments is None:
  115. queue.queue_arguments = {}
  116. self._set_max_priority(queue.queue_arguments)
  117. self[queue.name] = queue
  118. return queue
  119. def _set_ha_policy(self, args):
  120. policy = self.ha_policy
  121. if isinstance(policy, (list, tuple)):
  122. return args.update({'x-ha-policy': 'nodes',
  123. 'x-ha-policy-params': list(policy)})
  124. args['x-ha-policy'] = policy
  125. def _set_max_priority(self, args):
  126. if 'x-max-priority' not in args and self.max_priority is not None:
  127. return args.update({'x-max-priority': self.max_priority})
  128. def format(self, indent=0, indent_first=True):
  129. """Format routing table into string for log dumps."""
  130. active = self.consume_from
  131. if not active:
  132. return ''
  133. info = [QUEUE_FORMAT.strip().format(q)
  134. for _, q in sorted(items(active))]
  135. if indent_first:
  136. return textindent('\n'.join(info), indent)
  137. return info[0] + '\n' + textindent('\n'.join(info[1:]), indent)
  138. def select_add(self, queue, **kwargs):
  139. """Add new task queue that'll be consumed from.
  140. The queue will be active even when a subset has been selected
  141. using the :option:`celery worker -Q` option.
  142. """
  143. q = self.add(queue, **kwargs)
  144. if self._consume_from is not None:
  145. self._consume_from[q.name] = q
  146. return q
  147. def select(self, include):
  148. """Select a subset of currently defined queues to consume from.
  149. Arguments:
  150. include (Sequence[str], str): Names of queues to consume from.
  151. """
  152. if include:
  153. self._consume_from = {
  154. name: self[name] for name in maybe_list(include)
  155. }
  156. def deselect(self, exclude):
  157. """Deselect queues so that they won't be consumed from.
  158. Arguments:
  159. exclude (Sequence[str], str): Names of queues to avoid
  160. consuming from.
  161. """
  162. if exclude:
  163. exclude = maybe_list(exclude)
  164. if self._consume_from is None:
  165. # using all queues
  166. return self.select(k for k in self if k not in exclude)
  167. # using selection
  168. for queue in exclude:
  169. self._consume_from.pop(queue, None)
  170. def new_missing(self, name):
  171. return Queue(name, self.autoexchange(name), name)
  172. @property
  173. def consume_from(self):
  174. if self._consume_from is not None:
  175. return self._consume_from
  176. return self
  177. class AMQP(object):
  178. """App AMQP API: app.amqp."""
  179. Connection = Connection
  180. Consumer = Consumer
  181. Producer = Producer
  182. #: compat alias to Connection
  183. BrokerConnection = Connection
  184. queues_cls = Queues
  185. #: Cached and prepared routing table.
  186. _rtable = None
  187. #: Underlying producer pool instance automatically
  188. #: set by the :attr:`producer_pool`.
  189. _producer_pool = None
  190. # Exchange class/function used when defining automatic queues.
  191. # For example, you can use ``autoexchange = lambda n: None`` to use the
  192. # AMQP default exchange: a shortcut to bypass routing
  193. # and instead send directly to the queue named in the routing key.
  194. autoexchange = None
  195. #: Max size of positional argument representation used for
  196. #: logging purposes.
  197. argsrepr_maxsize = 1024
  198. #: Max size of keyword argument representation used for logging purposes.
  199. kwargsrepr_maxsize = 1024
  200. def __init__(self, app):
  201. self.app = app
  202. self.task_protocols = {
  203. 1: self.as_task_v1,
  204. 2: self.as_task_v2,
  205. }
  206. @cached_property
  207. def create_task_message(self):
  208. return self.task_protocols[self.app.conf.task_protocol]
  209. @cached_property
  210. def send_task_message(self):
  211. return self._create_task_sender()
  212. def Queues(self, queues, create_missing=None, ha_policy=None,
  213. autoexchange=None, max_priority=None):
  214. # Create new :class:`Queues` instance, using queue defaults
  215. # from the current configuration.
  216. conf = self.app.conf
  217. default_routing_key = conf.task_default_routing_key
  218. if create_missing is None:
  219. create_missing = conf.task_create_missing_queues
  220. if ha_policy is None:
  221. ha_policy = conf.task_queue_ha_policy
  222. if max_priority is None:
  223. max_priority = conf.task_queue_max_priority
  224. if not queues and conf.task_default_queue:
  225. queues = (Queue(conf.task_default_queue,
  226. exchange=self.default_exchange,
  227. routing_key=default_routing_key),)
  228. autoexchange = (self.autoexchange if autoexchange is None
  229. else autoexchange)
  230. return self.queues_cls(
  231. queues, self.default_exchange, create_missing,
  232. ha_policy, autoexchange, max_priority, default_routing_key,
  233. )
  234. def Router(self, queues=None, create_missing=None):
  235. """Return the current task router."""
  236. return _routes.Router(self.routes, queues or self.queues,
  237. self.app.either('task_create_missing_queues',
  238. create_missing), app=self.app)
  239. def flush_routes(self):
  240. self._rtable = _routes.prepare(self.app.conf.task_routes)
  241. def TaskConsumer(self, channel, queues=None, accept=None, **kw):
  242. if accept is None:
  243. accept = self.app.conf.accept_content
  244. return self.Consumer(
  245. channel, accept=accept,
  246. queues=queues or list(self.queues.consume_from.values()),
  247. **kw
  248. )
  249. def as_task_v2(self, task_id, name, args=None, kwargs=None,
  250. countdown=None, eta=None, group_id=None,
  251. expires=None, retries=0, chord=None,
  252. callbacks=None, errbacks=None, reply_to=None,
  253. time_limit=None, soft_time_limit=None,
  254. create_sent_event=False, root_id=None, parent_id=None,
  255. shadow=None, chain=None, now=None, timezone=None,
  256. origin=None, argsrepr=None, kwargsrepr=None):
  257. args = args or ()
  258. kwargs = kwargs or {}
  259. if not isinstance(args, (list, tuple)):
  260. raise TypeError('task args must be a list or tuple')
  261. if not isinstance(kwargs, Mapping):
  262. raise TypeError('task keyword arguments must be a mapping')
  263. if countdown: # convert countdown to ETA
  264. self._verify_seconds(countdown, 'countdown')
  265. now = now or self.app.now()
  266. timezone = timezone or self.app.timezone
  267. eta = maybe_make_aware(
  268. now + timedelta(seconds=countdown), tz=timezone,
  269. )
  270. if isinstance(expires, numbers.Real):
  271. self._verify_seconds(expires, 'expires')
  272. now = now or self.app.now()
  273. timezone = timezone or self.app.timezone
  274. expires = maybe_make_aware(
  275. now + timedelta(seconds=expires), tz=timezone,
  276. )
  277. eta = eta and eta.isoformat()
  278. # If we retry a task `expires` will already be ISO8601-formatted.
  279. if not isinstance(expires, string_t):
  280. expires = expires and expires.isoformat()
  281. if argsrepr is None:
  282. argsrepr = saferepr(args, self.argsrepr_maxsize)
  283. if kwargsrepr is None:
  284. kwargsrepr = saferepr(kwargs, self.kwargsrepr_maxsize)
  285. if JSON_NEEDS_UNICODE_KEYS: # pragma: no cover
  286. if callbacks:
  287. callbacks = [utf8dict(callback) for callback in callbacks]
  288. if errbacks:
  289. errbacks = [utf8dict(errback) for errback in errbacks]
  290. if chord:
  291. chord = utf8dict(chord)
  292. if not root_id: # empty root_id defaults to task_id
  293. root_id = task_id
  294. return task_message(
  295. headers={
  296. 'lang': 'py',
  297. 'task': name,
  298. 'id': task_id,
  299. 'shadow': shadow,
  300. 'eta': eta,
  301. 'expires': expires,
  302. 'group': group_id,
  303. 'retries': retries,
  304. 'timelimit': [time_limit, soft_time_limit],
  305. 'root_id': root_id,
  306. 'parent_id': parent_id,
  307. 'argsrepr': argsrepr,
  308. 'kwargsrepr': kwargsrepr,
  309. 'origin': origin or anon_nodename()
  310. },
  311. properties={
  312. 'correlation_id': task_id,
  313. 'reply_to': reply_to or '',
  314. },
  315. body=(
  316. args, kwargs, {
  317. 'callbacks': callbacks,
  318. 'errbacks': errbacks,
  319. 'chain': chain,
  320. 'chord': chord,
  321. },
  322. ),
  323. sent_event={
  324. 'uuid': task_id,
  325. 'root_id': root_id,
  326. 'parent_id': parent_id,
  327. 'name': name,
  328. 'args': argsrepr,
  329. 'kwargs': kwargsrepr,
  330. 'retries': retries,
  331. 'eta': eta,
  332. 'expires': expires,
  333. } if create_sent_event else None,
  334. )
  335. def as_task_v1(self, task_id, name, args=None, kwargs=None,
  336. countdown=None, eta=None, group_id=None,
  337. expires=None, retries=0,
  338. chord=None, callbacks=None, errbacks=None, reply_to=None,
  339. time_limit=None, soft_time_limit=None,
  340. create_sent_event=False, root_id=None, parent_id=None,
  341. shadow=None, now=None, timezone=None,
  342. **compat_kwargs):
  343. args = args or ()
  344. kwargs = kwargs or {}
  345. utc = self.utc
  346. if not isinstance(args, (list, tuple)):
  347. raise TypeError('task args must be a list or tuple')
  348. if not isinstance(kwargs, Mapping):
  349. raise TypeError('task keyword arguments must be a mapping')
  350. if countdown: # convert countdown to ETA
  351. self._verify_seconds(countdown, 'countdown')
  352. now = now or self.app.now()
  353. eta = now + timedelta(seconds=countdown)
  354. if isinstance(expires, numbers.Real):
  355. self._verify_seconds(expires, 'expires')
  356. now = now or self.app.now()
  357. expires = now + timedelta(seconds=expires)
  358. eta = eta and eta.isoformat()
  359. expires = expires and expires.isoformat()
  360. if JSON_NEEDS_UNICODE_KEYS: # pragma: no cover
  361. if callbacks:
  362. callbacks = [utf8dict(callback) for callback in callbacks]
  363. if errbacks:
  364. errbacks = [utf8dict(errback) for errback in errbacks]
  365. if chord:
  366. chord = utf8dict(chord)
  367. return task_message(
  368. headers={},
  369. properties={
  370. 'correlation_id': task_id,
  371. 'reply_to': reply_to or '',
  372. },
  373. body={
  374. 'task': name,
  375. 'id': task_id,
  376. 'args': args,
  377. 'kwargs': kwargs,
  378. 'group': group_id,
  379. 'retries': retries,
  380. 'eta': eta,
  381. 'expires': expires,
  382. 'utc': utc,
  383. 'callbacks': callbacks,
  384. 'errbacks': errbacks,
  385. 'timelimit': (time_limit, soft_time_limit),
  386. 'taskset': group_id,
  387. 'chord': chord,
  388. },
  389. sent_event={
  390. 'uuid': task_id,
  391. 'name': name,
  392. 'args': saferepr(args),
  393. 'kwargs': saferepr(kwargs),
  394. 'retries': retries,
  395. 'eta': eta,
  396. 'expires': expires,
  397. } if create_sent_event else None,
  398. )
  399. def _verify_seconds(self, s, what):
  400. if s < INT_MIN:
  401. raise ValueError('%s is out of range: %r' % (what, s))
  402. return s
  403. def _create_task_sender(self):
  404. default_retry = self.app.conf.task_publish_retry
  405. default_policy = self.app.conf.task_publish_retry_policy
  406. default_delivery_mode = self.app.conf.task_default_delivery_mode
  407. default_queue = self.default_queue
  408. queues = self.queues
  409. send_before_publish = signals.before_task_publish.send
  410. before_receivers = signals.before_task_publish.receivers
  411. send_after_publish = signals.after_task_publish.send
  412. after_receivers = signals.after_task_publish.receivers
  413. send_task_sent = signals.task_sent.send # XXX compat
  414. sent_receivers = signals.task_sent.receivers
  415. default_evd = self._event_dispatcher
  416. default_exchange = self.default_exchange
  417. default_rkey = self.app.conf.task_default_routing_key
  418. default_serializer = self.app.conf.task_serializer
  419. default_compressor = self.app.conf.result_compression
  420. def send_task_message(producer, name, message,
  421. exchange=None, routing_key=None, queue=None,
  422. event_dispatcher=None,
  423. retry=None, retry_policy=None,
  424. serializer=None, delivery_mode=None,
  425. compression=None, declare=None,
  426. headers=None, exchange_type=None, **kwargs):
  427. retry = default_retry if retry is None else retry
  428. headers2, properties, body, sent_event = message
  429. if headers:
  430. headers2.update(headers)
  431. if kwargs:
  432. properties.update(kwargs)
  433. qname = queue
  434. if queue is None and exchange is None:
  435. queue = default_queue
  436. if queue is not None:
  437. if isinstance(queue, string_t):
  438. qname, queue = queue, queues[queue]
  439. else:
  440. qname = queue.name
  441. if delivery_mode is None:
  442. try:
  443. delivery_mode = queue.exchange.delivery_mode
  444. except AttributeError:
  445. pass
  446. delivery_mode = delivery_mode or default_delivery_mode
  447. if exchange_type is None:
  448. try:
  449. exchange_type = queue.exchange.type
  450. except AttributeError:
  451. exchange_type = 'direct'
  452. # convert to anon-exchange, when exchange not set and direct ex.
  453. if (not exchange or not routing_key) and exchange_type == 'direct':
  454. exchange, routing_key = '', qname
  455. elif exchange is None:
  456. # not topic exchange, and exchange not undefined
  457. exchange = queue.exchange.name or default_exchange
  458. routing_key = routing_key or queue.routing_key or default_rkey
  459. if declare is None and queue and not isinstance(queue, Broadcast):
  460. declare = [queue]
  461. # merge default and custom policy
  462. retry = default_retry if retry is None else retry
  463. _rp = (dict(default_policy, **retry_policy) if retry_policy
  464. else default_policy)
  465. if before_receivers:
  466. send_before_publish(
  467. sender=name, body=body,
  468. exchange=exchange, routing_key=routing_key,
  469. declare=declare, headers=headers2,
  470. properties=properties, retry_policy=retry_policy,
  471. )
  472. ret = producer.publish(
  473. body,
  474. exchange=exchange,
  475. routing_key=routing_key,
  476. serializer=serializer or default_serializer,
  477. compression=compression or default_compressor,
  478. retry=retry, retry_policy=_rp,
  479. delivery_mode=delivery_mode, declare=declare,
  480. headers=headers2,
  481. **properties
  482. )
  483. if after_receivers:
  484. send_after_publish(sender=name, body=body, headers=headers2,
  485. exchange=exchange, routing_key=routing_key)
  486. if sent_receivers: # XXX deprecated
  487. if isinstance(body, tuple): # protocol version 2
  488. send_task_sent(
  489. sender=name, task_id=headers2['id'], task=name,
  490. args=body[0], kwargs=body[1],
  491. eta=headers2['eta'], taskset=headers2['group'],
  492. )
  493. else: # protocol version 1
  494. send_task_sent(
  495. sender=name, task_id=body['id'], task=name,
  496. args=body['args'], kwargs=body['kwargs'],
  497. eta=body['eta'], taskset=body['taskset'],
  498. )
  499. if sent_event:
  500. evd = event_dispatcher or default_evd
  501. exname = exchange
  502. if isinstance(exname, Exchange):
  503. exname = exname.name
  504. sent_event.update({
  505. 'queue': qname,
  506. 'exchange': exname,
  507. 'routing_key': routing_key,
  508. })
  509. evd.publish('task-sent', sent_event,
  510. producer, retry=retry, retry_policy=retry_policy)
  511. return ret
  512. return send_task_message
  513. @cached_property
  514. def default_queue(self):
  515. return self.queues[self.app.conf.task_default_queue]
  516. @cached_property
  517. def queues(self):
  518. """Queue name⇒ declaration mapping."""
  519. return self.Queues(self.app.conf.task_queues)
  520. @queues.setter # noqa
  521. def queues(self, queues):
  522. return self.Queues(queues)
  523. @property
  524. def routes(self):
  525. if self._rtable is None:
  526. self.flush_routes()
  527. return self._rtable
  528. @cached_property
  529. def router(self):
  530. return self.Router()
  531. @property
  532. def producer_pool(self):
  533. if self._producer_pool is None:
  534. self._producer_pool = pools.producers[
  535. self.app.connection_for_write()]
  536. self._producer_pool.limit = self.app.pool.limit
  537. return self._producer_pool
  538. publisher_pool = producer_pool # compat alias
  539. @cached_property
  540. def default_exchange(self):
  541. return Exchange(self.app.conf.task_default_exchange,
  542. self.app.conf.task_default_exchange_type)
  543. @cached_property
  544. def utc(self):
  545. return self.app.conf.enable_utc
  546. @cached_property
  547. def _event_dispatcher(self):
  548. # We call Dispatcher.publish with a custom producer
  549. # so don't need the diuspatcher to be enabled.
  550. return self.app.events.Dispatcher(enabled=False)