Source

amqpev / amqpev / api.py

Full commit
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
# Copyright (c) 2009 Kyle Schaffrick
#
# This module is part of AMQPev and is released under the MIT license. See
# COPYING in the root of this distribution for details.

import datetime
from protocol import Connection
from transport import GreenTCPTransport
from spec.spec import spec_0_8 as spec
import eventlet.api as ev_api

import logging

log = logging.getLogger("amqpev.api")

MESSAGE_CONTENT_TYPES = {}

class BrokerServiceFactory(object):
    """
    This class is a convenience method for setting up an AMQP transport,
    connection, and creating broker service objects and their underlying
    channels.
    """
    def __init__(self, host):
        """
        Create a new broker service factory. The constructor will create a new
        AMQP connection to the specified host[:port] and start it immediately.
        """
        trans = GreenTCPTransport(host)
        self.conn = Connection(trans)

    def __call__(self):
        """
        Create a new channel wrapped in a broker service object.
        """
        return BrokerService(self.conn.channel())

    def close(self):
        """
        Close the connection associated with this broker service factory.
        """
        self.conn.close()


class BrokerService(object):
    """
    A BrokerService object along with it's helper classes for the various
    objects in AMQP model provide a high-level API for interacting with the
    AMQP broker services. A BrokerService is essentially an adapter around a
    pre-configured channel. It implements the context manager interface for use
    with the with statement, which will automatically manage the lifetime of
    the underlying channel.
    """
    def __init__(self, channel):
        self._ch = channel
        self._exchanges = {
                '': Exchange(self._ch, "", 'direct', False, False, {},
                    pre_declared=True),
                'amq.direct': Exchange(self._ch, "amq.direct", 'direct', False,
                    False, {}, pre_declared=True),
                'amq.fanout': Exchange(self._ch, "amq.fanout", 'fanout', False,
                    False, {}, pre_declared=True),
                'amq.topic': Exchange(self._ch, "amq.topic", 'topic', False,
                    False, {}, pre_declared=True) }
        self._queues = {}

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        try:
            self.close()
        except:
            if type is None:
                raise
            else:
                log.exception("Masking exception while shutting down"
                        " BrokerService")

    def close(self):
        """
        Close the broker service's channel.
        """
        self._ch.close()

    @property
    def channel(self):
        """
        Return the channel associated with this broker service.
        """
        return self._ch

    def exchange(self, name=None, type='direct', passive=False, durable=False,
            arguments={}):
        """
        Create an exchange. If the specified exchange has been seen before,
        return the same Exchange object from our identity map. If we haven't,
        declare it on the channel, save, and return the new Exchange.
        """
        if name is None:
            raise ValueError("Exchanges must be explicitly named")
        elif name in self._exchanges:
            return self._exchanges[name]
        else:
            ex = Exchange(self._ch, name, type, passive, durable, arguments)
            self._exchanges[ex.name] = ex
            return ex

    def queue(self, name=None, passive=False, durable=False, exclusive=None,
            auto_delete=None, arguments=None):
        """
        Create a queue. If we've seen this queue name before, return the same
        Queue object from our identity map. If we haven't, declare it on the
        channel, save, and return.

        If no name is given, a new, exclusive, auto-delete queue will always be
        declared and returned. However, if a name is given, auto_delete
        defaults to False.
        """
        if exclusive is None:
            exclusive = name is None
        if auto_delete is None:
            auto_delete = name is None
        arguments = arguments or {}

        if name is not None and name in self._queues:
            return self._queues[name]
        else:
            q = Queue(self._ch, self, name, passive, durable, exclusive,
                    auto_delete, arguments)
            self._queues[q.name] = q
            return q

    def publish(self, message, routing_key, **kwargs):
        """
        Convenience for publishing on the default direct exchange, to which all
        queues are implicitly bound by name. This implicit binding uses the
        queue name as the routing key, therefore a message may be queued
        directly in any queue by sending to the default exchange using the
        queue's name as the routing key.
        """
        return self.exchange('').publish(message, routing_key, **kwargs)

    def consume_from(self, *queues, **kwargs):
        """
        Consume from some set of queues. Returns a context manager object,
        specifically a ConsumerSet.
        """
        return ConsumerSet(self._ch, queues, **kwargs)

    def ack(self, message, thru=False):
        """
        Acknowledge a message.
        """
        self._ch.send_method(spec.basic.ack,
                delivery_tag=message.delivery_headers['delivery_tag'],
                multiple=thru)


class Exchange(object):
    """
    This is a helper object representing an AMQP exchange in the high-level
    API.
    """
    def __init__(self, channel, name, type, passive, durable, arguments,
            pre_declared=False):
        self._ch = channel
        self._name = name

        if not pre_declared:
            self._ch.sync_method(spec.exchange.declare,
                    (spec.exchange.declare_ok,), ticket=0, exchange=name,
                    type=type, passive=passive, durable=durable,
                    auto_delete=False, internal=False, nowait=False,
                    arguments=arguments)

    @property
    def name(self):
        """
        The name of the exchange as declared on the broker.
        """
        return self._name

    def publish(self, message, routing_key="", mandatory=False, immediate=False):
        """
        Publish a message object to this exchange, optionally with the
        specified routing key. By default, the mandatory and immediate delivery
        flags are not set.
        """
        self._ch.send_method(spec.basic.publish, ticket=0, exchange=self._name,
                routing_key=routing_key, mandatory=mandatory,
                immediate=immediate)
        self._ch.send_content(spec.basic, message.properties, message.encode(),
                len(message.body))


class Queue(object):
    """
    This helper object represents an AMQP queue in the high-level API.
    """
    def __init__(self, channel, broker_svc, name, passive, durable, exclusive,
            auto_delete, arguments):
        self._ch = channel
        self._bs = broker_svc
        self._name = "" if name is None else name

        rx_m = self._ch.sync_method(spec.queue.declare,
                (spec.queue.declare_ok,), ticket=0, queue=self._name,
                passive=passive, durable=durable, exclusive=exclusive,
                auto_delete=auto_delete, nowait=False, arguments=arguments)

        self._name = rx_m.args['queue']

    @property
    def name(self):
        """
        The name of the queue as declared on the broker.
        """
        return self._name

    def bind(self, exch, routing_key="", arguments=None):
        """
        Bind this queue to the given exchange object, optionally with the
        specified routing key and binding arguments.
        """
        arguments = arguments or {}
        self._ch.sync_method(spec.queue.bind, (spec.queue.bind_ok,), ticket=0,
                queue=self.name, exchange=exch.name, routing_key=routing_key,
                nowait=False, arguments=arguments)
        return self

    def publish(self, message, **kwargs):
        """
        Convenience for publishing to this queue via the default direct
        exchange to which all queues are implicitly bound.
        """
        return self._bs.publish(message, self.name, **kwargs)

    def get(self, no_ack=True):
        """
        Perform a synchronous get operation on this queue. Will return a
        message or None if the queue is empty. The no_ack flag's default of
        True instructs the broker to automatically acknowledge the message as
        delivered. When set to False, the message may be re-delivered by the
        broker if it does not receive an explicit acknowledgement.
        """
        rx_m = self._ch.sync_method(spec.basic.get,
                (spec.basic.get_ok, spec.basic.get_empty),
                ticket=0, queue=self.name, no_ack=no_ack)

        if rx_m.method == spec.basic.get_empty:
            return None

        (properties, body_len, body) = self._ch.recv_content()
        msg = BaseMessage.create_decoded(body, properties)
        msg.delivery_headers = rx_m.args.copy()
        msg._consumed_from_queue = self
        return msg

    def __contains__(self, message):
        """
        This magic method allows neatly testing if a delivered message came
        from a particular queue using the expression "msg in queue". Returns
        True if the message was delivered from this queue.
        """
        return message._consumed_from_queue is self


class BaseMessage(object):
    """
    An AMQP message, with payload, properties, and delivery headers. This class
    is an abstract base class, subclasses are used to represent messages that
    are delivered with particular content types.

    Subclasses of this message should implement an instance method encode()
    which returns the encoded data that will be inserted into the body of the
    AMQP message when it is sent. They should also implement a class method
    decode(encoded_body, properties) that constructs an instance given the
    encoded body and properties received from the AMQP message.
    """
    def __init__(self, body, properties):
        self.body = body
        self.properties = properties
        self.delivery_headers = {}

    @staticmethod
    def add_content_type(content_type, class_):
        """
        Register a subclass of Message that will be used to represent messages
        sent with the specified content type property.
        """
        MESSAGE_CONTENT_TYPES[content_type] = class_

    @staticmethod
    def create_decoded(body, properties):
        """
        Find the appropriate message class for the content type, and construct
        the message.
        """
        content_type = properties.get('content_type', 'text/plain')
        class_ = MESSAGE_CONTENT_TYPES.get(content_type, Message)
        return class_.decode(body, properties)


class Message(BaseMessage):
    """
    An AMQP message which sets some commonly used properties to sensible
    defaults. It performs no encoding or decoding on its body, sets the
    content_type to a default of 'text/plain', and sets the timestamp to the
    current UTC date/time.
    """
    def __init__(self, body, **kwargs):
        effective_props = {
                'content_type': 'text/plain',
                'timestamp': datetime.datetime.utcnow() }
        effective_props.update(kwargs)
        BaseMessage.__init__(self, body, effective_props)

    def encode(self):
        return self.body

    @classmethod
    def decode(cls, encoded_body, properties):
        return cls(encoded_body, **properties)


class ConsumerSet(object):
    """
    A set of consumers that can be used to asynchronously consume messages from
    some set of queues. This object implements the context manager interface,
    which can be used to automatically manage the lifetime of the underlying
    consumers. This class also implements the iterator protocol for consuming
    messages asynchronously.

    It should not be necessary to construct a ConsumerSet directly, as
    BrokerService objects provide a consume_from() factory method for producing
    them.
    """
    def __init__(self, channel, queues, no_ack=False, prefetch=None,
            timeout=None, raise_on_timeout=False):
        self._ch = channel
        self.queues = queues
        self._c_tags = {}
        self._state = 'stopped'
        self._no_ack = no_ack
        self._prefetch = prefetch
        self._raise_on_timeout = raise_on_timeout
        self._timeout = timeout

    def __enter__(self):
        self.start_consumers()
        return self

    def __exit__(self, type, value, tb):
        try:
            self.stop_consumers()
            # This is to consume any pending messages from the receive queues.
            # stop_consumers() does not do this because the application may
            # call stop_consumers() and continue iterating to actually consume
            # those messages intead of discarding them.
            list(self)
        except:
            if type is None:
                raise
            else:
                log.exception("Masking exception while shutting down"
                        " ConsumerSet.")

    def start_consumers(self):
        """
        Start consumers with the broker on the specified queues. This is
        automatically called on entry when the ConsumerSet is used in a with
        block.
        """
        if self._state == 'running': return

        if self._prefetch:
            self._ch.sync_method(spec.basic.qos, (spec.basic.qos_ok,),
                    prefetch_size=0, prefetch_count=self._prefetch,
                    global_=False)

        for queue in self.queues:
            rx_m = self._ch.sync_method(spec.basic.consume,
                    (spec.basic.consume_ok,), ticket=0, queue=queue.name,
                    consumer_tag='', no_local=False, no_ack=self._no_ack,
                    exclusive=False, nowait=False)
            self._c_tags[ rx_m.args['consumer_tag'] ] = queue

        self._state = 'running'

    def stop_consumers(self):
        """
        Stop the consumers on the broker. This is automatically called on exit
        when the ConsumerSet is used in a with block.

        Note that calling stop_consumers() may not stop the flow of messages
        immediately. Any number of messages may be emitted by the ConsumerSet
        after this call occurs, especially with a high (or no) prefetch limit.
        Such messages include deliveries sent during the network and buffering
        latency between the client and the broker, as well as any unseen
        deliveries that are still in transit, or queued in the OS's socket
        receive queue or AMQPev's channel demux queues.

        See next_message() and __iter__() docstrings for more details on this
        behavior.
        """
        if self._state != 'running': return 

        for c_tag in self._c_tags.keys():
            self._ch.send_method(spec.basic.cancel, consumer_tag=c_tag,
                    nowait=False)

        self._state = 'stop-wait'

    def _next_message(self):
        """
        Return the next delivered message. See next_message() and __iter__()'s
        docstrings for more information, as they are the public API for this
        implementation method.

        The main difference is that if a timeout is specified, this
        implementation always raises TimeoutError. Honoring raise_on_timeout is
        done in the __iter__ and next_message() interfaces.
        """
        if self._state == 'stopped':
            return

        if self._timeout is not None:
            rx_m = ev_api.with_timeout(self._timeout, self._ch.recv_method)
        else:
            rx_m = self._ch.recv_method()

        if rx_m.method == spec.basic.deliver:
            (properties, body_len, body) = self._ch.recv_content()
            msg = BaseMessage.create_decoded(body, properties)
            msg.delivery_headers = rx_m.args.copy()
            msg._consumed_from_queue = self._c_tags[
                    msg.delivery_headers['consumer_tag'] ]
            return msg

        elif rx_m.method == spec.basic.cancel_ok:
            self._c_tags.pop(rx_m.args['consumer_tag'])

            if not self._c_tags:
                if self._prefetch:
                    self._ch.sync_method(spec.basic.qos, (spec.basic.qos_ok,),
                            prefetch_size=0, prefetch_count=0, global_=False)
                self._state == 'stopped'
                return

    def next_message(self):
        """
        Return the next message delivered. This call will block the current
        coroutine until a message is received, or (if timeout was specified
        when constructing the ConsumerSet) the specified timeout elapses.

        The method will return None immediately when called after a
        stop_consumers() call, and all the messages the broker delivered have
        been returned.

        The raise_on_timeout argument specifies behavior when a timeout occurs.
        When it is False, next_message() will simply return None, which is the
        default. If it is True, a TimeoutError is raised. This can be used if
        the application needs to differentiate between a timeout condition and
        a no-more-deliveries condition.
        """
        try:
            return self._next_message()
        except ev_api.TimeoutError:
            if self._raise_on_error: raise
            else: return None

    def __iter__(self):
        """
        Iterate through delivered messages. The iterator will block the current
        coroutine until a message is received, or (if timeout was specified
        when constructing the ConsumerSet) the specified timeout elapses.

        The raise_on_timeout argument specifies behavior when a timeout occurs.
        When it is False, the iterator will simply repeatedly yield None each
        time the timeout period elapses without receiving a message. When it is
        True, a TimeoutError will be raised.

        In the case that raise_on_timeout is False or no timeout is specified,
        iteration will not end (either repeatedly yielding None or blocking
        indefinitely) until stop_consumers() is called. Any number of messages
        may be yielded by the iterator following the stop_consumers() call, but
        iteration will end when the broker confirms it has stopped delivery (in
        AMQP terms, when we receive a basic.cancel_ok method for each
        consumer). Thus, when the iteration ends, it can be assumed that all
        messages the broker delivered have been emitted by the iterator.
        """
        while True:
            try:
                msg = self._next_message()
                if msg is None: return
                else: yield msg
            except ev_api.TimeoutError:
                if self._raise_on_error: raise
                else: yield None


try:
    from jsonmessage import JSONMessage
except ImportError:
    pass