PageRenderTime 68ms CodeModel.GetById 32ms RepoModel.GetById 0ms app.codeStats 0ms

/carrot/backends/pyamqplib.py

https://bitbucket.org/invenshure/carrot
Python | 364 lines | 300 code | 33 blank | 31 comment | 27 complexity | 38a8ffb90d69f6ee2e4f64ae19299d40 MD5 | raw file
  1. """
  2. `amqplib`_ backend for carrot.
  3. .. _`amqplib`: http://barryp.org/software/py-amqplib/
  4. """
  5. from amqplib.client_0_8 import transport
  6. # amqplib's handshake mistakenly identifies as protocol version 1191,
  7. # this breaks in RabbitMQ tip, which no longer falls back to
  8. # 0-8 for unknown ids.
  9. transport.AMQP_PROTOCOL_HEADER = "AMQP\x01\x01\x08\x00"
  10. from amqplib import client_0_8 as amqp
  11. from amqplib.client_0_8.exceptions import AMQPConnectionException
  12. from amqplib.client_0_8.exceptions import AMQPChannelException
  13. from amqplib.client_0_8.serialization import AMQPReader, AMQPWriter
  14. from carrot.backends.base import BaseMessage, BaseBackend
  15. from itertools import count
  16. import socket
  17. import warnings
  18. import weakref
  19. DEFAULT_PORT = 5672
  20. class Connection(amqp.Connection):
  21. def drain_events(self, allowed_methods=None, timeout=None):
  22. """Wait for an event on any channel."""
  23. return self.wait_multi(self.channels.values(), timeout=timeout)
  24. def wait_multi(self, channels, allowed_methods=None, timeout=None):
  25. """Wait for an event on a channel."""
  26. chanmap = dict((chan.channel_id, chan) for chan in channels)
  27. chanid, method_sig, args, content = self._wait_multiple(
  28. chanmap.keys(), allowed_methods, timeout=timeout)
  29. channel = chanmap[chanid]
  30. if content \
  31. and channel.auto_decode \
  32. and hasattr(content, 'content_encoding'):
  33. try:
  34. content.body = content.body.decode(content.content_encoding)
  35. except Exception:
  36. pass
  37. amqp_method = channel._METHOD_MAP.get(method_sig, None)
  38. if amqp_method is None:
  39. raise Exception('Unknown AMQP method (%d, %d)' % method_sig)
  40. if content is None:
  41. return amqp_method(channel, args)
  42. else:
  43. return amqp_method(channel, args, content)
  44. def read_timeout(self, timeout=None):
  45. if timeout is None:
  46. return self.method_reader.read_method()
  47. sock = self.transport.sock
  48. prev = sock.gettimeout()
  49. sock.settimeout(timeout)
  50. try:
  51. return self.method_reader.read_method()
  52. finally:
  53. sock.settimeout(prev)
  54. def _wait_multiple(self, channel_ids, allowed_methods, timeout=None):
  55. for channel_id in channel_ids:
  56. method_queue = self.channels[channel_id].method_queue
  57. for queued_method in method_queue:
  58. method_sig = queued_method[0]
  59. if (allowed_methods is None) \
  60. or (method_sig in allowed_methods) \
  61. or (method_sig == (20, 40)):
  62. method_queue.remove(queued_method)
  63. method_sig, args, content = queued_method
  64. return channel_id, method_sig, args, content
  65. # Nothing queued, need to wait for a method from the peer
  66. while True:
  67. channel, method_sig, args, content = self.read_timeout(timeout)
  68. if (channel in channel_ids) \
  69. and ((allowed_methods is None) \
  70. or (method_sig in allowed_methods) \
  71. or (method_sig == (20, 40))):
  72. return channel, method_sig, args, content
  73. # Not the channel and/or method we were looking for. Queue
  74. # this method for later
  75. self.channels[channel].method_queue.append((method_sig,
  76. args,
  77. content))
  78. #
  79. # If we just queued up a method for channel 0 (the Connection
  80. # itself) it's probably a close method in reaction to some
  81. # error, so deal with it right away.
  82. #
  83. if channel == 0:
  84. self.wait()
  85. class QueueAlreadyExistsWarning(UserWarning):
  86. """A queue with that name already exists, so a recently changed
  87. ``routing_key`` or other settings might be ignored unless you
  88. rename the queue or restart the broker."""
  89. class Message(BaseMessage):
  90. """A message received by the broker.
  91. Usually you don't insantiate message objects yourself, but receive
  92. them using a :class:`carrot.messaging.Consumer`.
  93. :param backend: see :attr:`backend`.
  94. :param amqp_message: see :attr:`_amqp_message`.
  95. .. attribute:: body
  96. The message body.
  97. .. attribute:: delivery_tag
  98. The message delivery tag, uniquely identifying this message.
  99. .. attribute:: backend
  100. The message backend used.
  101. A subclass of :class:`carrot.backends.base.BaseBackend`.
  102. .. attribute:: _amqp_message
  103. A :class:`amqplib.client_0_8.basic_message.Message` instance.
  104. This is a private attribute and should not be accessed by
  105. production code.
  106. """
  107. def __init__(self, backend, amqp_message, **kwargs):
  108. self._amqp_message = amqp_message
  109. self.backend = backend
  110. for attr_name in ("body",
  111. "delivery_tag",
  112. "content_type",
  113. "content_encoding",
  114. "delivery_info"):
  115. kwargs[attr_name] = getattr(amqp_message, attr_name, None)
  116. super(Message, self).__init__(backend, **kwargs)
  117. class Backend(BaseBackend):
  118. """amqplib backend
  119. :param connection: see :attr:`connection`.
  120. .. attribute:: connection
  121. A :class:`carrot.connection.BrokerConnection` instance. An established
  122. connection to the broker.
  123. """
  124. default_port = DEFAULT_PORT
  125. connection_errors = (AMQPConnectionException,
  126. socket.error,
  127. IOError,
  128. OSError)
  129. channel_errors = (AMQPChannelException, )
  130. Message = Message
  131. def __init__(self, connection, **kwargs):
  132. self.connection = connection
  133. self.default_port = kwargs.get("default_port", self.default_port)
  134. self._channel_ref = None
  135. @property
  136. def _channel(self):
  137. return callable(self._channel_ref) and self._channel_ref()
  138. @property
  139. def channel(self):
  140. """If no channel exists, a new one is requested."""
  141. if not self._channel:
  142. connection = self.connection.connection
  143. self._channel_ref = weakref.ref(connection.channel())
  144. return self._channel
  145. def establish_connection(self):
  146. """Establish connection to the AMQP broker."""
  147. conninfo = self.connection
  148. if not conninfo.hostname:
  149. raise KeyError("Missing hostname for AMQP connection.")
  150. if conninfo.userid is None:
  151. raise KeyError("Missing user id for AMQP connection.")
  152. if conninfo.password is None:
  153. raise KeyError("Missing password for AMQP connection.")
  154. if not conninfo.port:
  155. conninfo.port = self.default_port
  156. return Connection(host=conninfo.host,
  157. userid=conninfo.userid,
  158. password=conninfo.password,
  159. virtual_host=conninfo.virtual_host,
  160. insist=conninfo.insist,
  161. ssl=conninfo.ssl,
  162. connect_timeout=conninfo.connect_timeout)
  163. def close_connection(self, connection):
  164. """Close the AMQP broker connection."""
  165. connection.close()
  166. def queue_exists(self, queue):
  167. """Check if a queue has been declared.
  168. :rtype bool:
  169. """
  170. try:
  171. self.channel.queue_declare(queue=queue, passive=True)
  172. except AMQPChannelException, e:
  173. if e.amqp_reply_code == 404:
  174. return False
  175. raise e
  176. else:
  177. return True
  178. def queue_delete(self, queue, if_unused=False, if_empty=False):
  179. """Delete queue by name."""
  180. return self.channel.queue_delete(queue, if_unused, if_empty)
  181. def queue_purge(self, queue, **kwargs):
  182. """Discard all messages in the queue. This will delete the messages
  183. and results in an empty queue."""
  184. return self.channel.queue_purge(queue=queue)
  185. def queue_declare(self, queue, durable, exclusive, auto_delete,
  186. warn_if_exists=False, arguments=None):
  187. """Declare a named queue."""
  188. if warn_if_exists and self.queue_exists(queue):
  189. warnings.warn(QueueAlreadyExistsWarning(
  190. QueueAlreadyExistsWarning.__doc__))
  191. return self.channel.queue_declare(queue=queue,
  192. durable=durable,
  193. exclusive=exclusive,
  194. auto_delete=auto_delete,
  195. arguments=arguments)
  196. def exchange_declare(self, exchange, type, durable, auto_delete):
  197. """Declare an named exchange."""
  198. return self.channel.exchange_declare(exchange=exchange,
  199. type=type,
  200. durable=durable,
  201. auto_delete=auto_delete)
  202. def queue_bind(self, queue, exchange, routing_key, arguments=None):
  203. """Bind queue to an exchange using a routing key."""
  204. return self.channel.queue_bind(queue=queue,
  205. exchange=exchange,
  206. routing_key=routing_key,
  207. arguments=arguments)
  208. def message_to_python(self, raw_message):
  209. """Convert encoded message body back to a Python value."""
  210. return self.Message(backend=self, amqp_message=raw_message)
  211. def get(self, queue, no_ack=False):
  212. """Receive a message from a declared queue by name.
  213. :returns: A :class:`Message` object if a message was received,
  214. ``None`` otherwise. If ``None`` was returned, it probably means
  215. there was no messages waiting on the queue.
  216. """
  217. raw_message = self.channel.basic_get(queue, no_ack=no_ack)
  218. if not raw_message:
  219. return None
  220. return self.message_to_python(raw_message)
  221. def declare_consumer(self, queue, no_ack, callback, consumer_tag,
  222. nowait=False):
  223. """Declare a consumer."""
  224. return self.channel.basic_consume(queue=queue,
  225. no_ack=no_ack,
  226. callback=callback,
  227. consumer_tag=consumer_tag,
  228. nowait=nowait)
  229. def consume(self, limit=None):
  230. """Returns an iterator that waits for one message at a time."""
  231. for total_message_count in count():
  232. if limit and total_message_count >= limit:
  233. raise StopIteration
  234. if not self.channel.is_open:
  235. raise StopIteration
  236. self.channel.wait()
  237. yield True
  238. def cancel(self, consumer_tag):
  239. """Cancel a channel by consumer tag."""
  240. if not self.channel.connection:
  241. return
  242. self.channel.basic_cancel(consumer_tag)
  243. def close(self):
  244. """Close the channel if open."""
  245. if self._channel and self._channel.is_open:
  246. self._channel.close()
  247. self._channel_ref = None
  248. def ack(self, delivery_tag):
  249. """Acknowledge a message by delivery tag."""
  250. return self.channel.basic_ack(delivery_tag)
  251. def reject(self, delivery_tag):
  252. """Reject a message by deliver tag."""
  253. return self.channel.basic_reject(delivery_tag, requeue=False)
  254. def requeue(self, delivery_tag):
  255. """Reject and requeue a message by delivery tag."""
  256. return self.channel.basic_reject(delivery_tag, requeue=True)
  257. def prepare_message(self, message_data, delivery_mode, priority=None,
  258. content_type=None, content_encoding=None):
  259. """Encapsulate data into a AMQP message."""
  260. message = amqp.Message(message_data, priority=priority,
  261. content_type=content_type,
  262. content_encoding=content_encoding)
  263. message.properties["delivery_mode"] = delivery_mode
  264. return message
  265. def publish(self, message, exchange, routing_key, mandatory=None,
  266. immediate=None, headers=None):
  267. """Publish a message to a named exchange."""
  268. if headers:
  269. message.properties["headers"] = headers
  270. ret = self.channel.basic_publish(message, exchange=exchange,
  271. routing_key=routing_key,
  272. mandatory=mandatory,
  273. immediate=immediate)
  274. if mandatory or immediate:
  275. self.close()
  276. def qos(self, prefetch_size, prefetch_count, apply_global=False):
  277. """Request specific Quality of Service."""
  278. self.channel.basic_qos(prefetch_size, prefetch_count,
  279. apply_global)
  280. def flow(self, active):
  281. """Enable/disable flow from peer."""
  282. self.channel.flow(active)