/carrot/backends/pyamqplib.py

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