celery/celery/app/amqp.py

# -*- coding: utf-8 -*-
"""
    celery.app.amqp
    ~~~~~~~~~~~~~~~

    Sending and receiving messages using Kombu.

"""
from __future__ import absolute_import

from datetime import timedelta
from weakref import WeakValueDictionary

from kombu import Connection, Consumer, Exchange, Producer, Queue
from kombu.common import entry_to_queue
from kombu.pools import ProducerPool
from kombu.utils.encoding import safe_repr

from celery import signals
from celery.utils import cached_property, uuid
from celery.utils.text import indent as textindent

from . import app_or_default
from . import routes as _routes

#: Human readable queue declaration.
QUEUE_FORMAT = """
. %(name)s exchange:%(exchange)s(%(exchange_type)s) binding:%(routing_key)s
"""


class Queues(dict):
    """Queue name⇒ declaration mapping.

    :param queues: Initial list/tuple or dict of queues.
    :keyword create_missing: By default any unknown queues will be
                             added automatically, but if disabled
                             the occurrence of unknown queues
                             in `wanted` will raise :exc:`KeyError`.


    """
    #: If set, this is a subset of queues to consume from.
    #: The rest of the queues are then used for routing only.
    _consume_from = None

    def __init__(self, queues=None, default_exchange=None,
            create_missing=True):
        dict.__init__(self)
        self.aliases = WeakValueDictionary()
        self.default_exchange = default_exchange
        self.create_missing = create_missing
        if isinstance(queues, (tuple, list)):
            queues = dict((q.name, q) for q in queues)
        for name, q in (queues or {}).iteritems():
            self.add(q) if isinstance(q, Queue) else self.add_compat(name, **q)

    def __getitem__(self, name):
        try:
            return self.aliases[name]
        except KeyError:
            return dict.__getitem__(self, name)

    def __setitem__(self, name, queue):
        if self.default_exchange and (not queue.exchange or
                                      not queue.exchange.name):
            queue.exchange = self.default_exchange
        dict.__setitem__(self, name, queue)
        if queue.alias:
            self.aliases[queue.alias] = queue

    def __missing__(self, name):
        if self.create_missing:
            return self.add(self.new_missing(name))
        raise KeyError(name)

    def add(self, queue, **kwargs):
        """Add new queue.

        :param queue: Name of the queue.
        :keyword exchange: Name of the exchange.
        :keyword routing_key: Binding key.
        :keyword exchange_type: Type of exchange.
        :keyword \*\*options: Additional declaration options.

        """
        if not isinstance(queue, Queue):
            return self.add_compat(queue, **kwargs)
        self[queue.name] = queue
        return queue

    def add_compat(self, name, **options):
        # docs used to use binding_key as routing key
        options.setdefault('routing_key', options.get('binding_key'))
        q = self[name] = entry_to_queue(name, **options)
        return q

    def format(self, indent=0, indent_first=True):
        """Format routing table into string for log dumps."""
        active = self.consume_from
        if not active:
            return ''
        info = [QUEUE_FORMAT.strip() % {
                    'name': (name + ':').ljust(12),
                    'exchange': q.exchange.name,
                    'exchange_type': q.exchange.type,
                    'routing_key': q.routing_key}
                        for name, q in sorted(active.iteritems())]
        if indent_first:
            return textindent('\n'.join(info), indent)
        return info[0] + '\n' + textindent('\n'.join(info[1:]), indent)

    def select_subset(self, wanted):
        """Sets :attr:`consume_from` by selecting a subset of the
        currently defined queues.

        :param wanted: List of wanted queue names.
        """
        if wanted:
            self._consume_from = dict((name, self[name]) for name in wanted)

    def select_remove(self, queue):
        if self._consume_from is None:
            self.select_subset(k for k in self.keys() if k != queue)
        else:
            self._consume_from.pop(queue, None)

    def new_missing(self, name):
        return Queue(name, Exchange(name), name)

    @property
    def consume_from(self):
        if self._consume_from is not None:
            return self._consume_from
        return self


class TaskProducer(Producer):
    app = None
    auto_declare = False
    retry = False
    retry_policy = None

    def __init__(self, channel=None, exchange=None, *args, **kwargs):
        self.retry = kwargs.pop('retry', self.retry)
        self.retry_policy = kwargs.pop('retry_policy',
                                        self.retry_policy or {})
        exchange = exchange or self.exchange
        self.queues = self.app.amqp.queues  # shortcut
        super(TaskProducer, self).__init__(channel, exchange, *args, **kwargs)

    def publish_task(self, task_name, task_args=None, task_kwargs=None,
            countdown=None, eta=None, task_id=None, group_id=None,
            taskset_id=None,  # compat alias to group_id
            expires=None, exchange=None, exchange_type=None,
            event_dispatcher=None, retry=None, retry_policy=None,
            queue=None, now=None, retries=0, chord=None, callbacks=None,
            errbacks=None, mandatory=None, priority=None, immediate=None,
            routing_key=None, serializer=None, delivery_mode=None,
            compression=None, timeout=None, soft_timeout=None,
            **kwargs):
        """Send task message."""
        # merge default and custom policy
        retry = self.retry if retry is None else retry
        _rp = (dict(self.retry_policy, **retry_policy) if retry_policy
                                                       else self.retry_policy)
        task_id = task_id or uuid()
        task_args = task_args or []
        task_kwargs = task_kwargs or {}
        if not isinstance(task_args, (list, tuple)):
            raise ValueError('task args must be a list or tuple')
        if not isinstance(task_kwargs, dict):
            raise ValueError('task kwargs must be a dictionary')
        if countdown:  # Convert countdown to ETA.
            now = now or self.app.now()
            eta = now + timedelta(seconds=countdown)
        if isinstance(expires, (int, float)):
            now = now or self.app.now()
            expires = now + timedelta(seconds=expires)
        eta = eta and eta.isoformat()
        expires = expires and expires.isoformat()
        timeouts = (timeout, soft_timeout)

        body = {'task': task_name,
                'id': task_id,
                'args': task_args,
                'kwargs': task_kwargs,
                'retries': retries or 0,
                'eta': eta,
                'expires': expires,
                'utc': self.utc,
                'callbacks': callbacks,
                'errbacks': errbacks,
                'timeouts': timeouts}
        group_id = group_id or taskset_id
        if group_id:
            body['taskset'] = group_id
        if chord:
            body['chord'] = chord

        self.publish(body, exchange=exchange, mandatory=mandatory,
             immediate=immediate, routing_key=routing_key,
             serializer=serializer or self.serializer,
             compression=compression or self.compression,
             retry=retry, retry_policy=_rp, delivery_mode=delivery_mode,
             priority=priority, declare=[self.queues[queue]] if queue else [],
             **kwargs)

        signals.task_sent.send(sender=task_name, **body)
        if event_dispatcher:
            exname = exchange or self.exchange
            if isinstance(exname, Exchange):
                exname = exname.name
            event_dispatcher.send('task-sent', uuid=task_id,
                                               name=task_name,
                                               args=safe_repr(task_args),
                                               kwargs=safe_repr(task_kwargs),
                                               retries=retries,
                                               eta=eta,
                                               expires=expires,
                                               queue=queue,
                                               exchange=exname,
                                               routing_key=routing_key)
        return task_id
    delay_task = publish_task   # XXX Compat


class TaskPublisher(TaskProducer):
    """Deprecated version of :class:`TaskProducer`."""

    def __init__(self, channel=None, exchange=None, *args, **kwargs):
        self.app = app_or_default(kwargs.pop('app', self.app))
        self.retry = kwargs.pop('retry', self.retry)
        self.retry_policy = kwargs.pop('retry_policy',
                                        self.retry_policy or {})
        exchange = exchange or self.exchange
        if not isinstance(exchange, Exchange):
            exchange = Exchange(exchange,
                                kwargs.pop('exchange_type', 'direct'))
        self.queues = self.app.amqp.queues  # shortcut
        super(TaskPublisher, self).__init__(channel, exchange, *args, **kwargs)


class TaskConsumer(Consumer):
    app = None

    def __init__(self, channel, queues=None, app=None, **kw):
        self.app = app or self.app
        super(TaskConsumer, self).__init__(channel,
                queues or self.app.amqp.queues.consume_from.values(), **kw)


class AMQP(object):
    Connection = Connection
    Consumer = Consumer

    #: compat alias to Connection
    BrokerConnection = Connection

    #: Cached and prepared routing table.
    _rtable = None

    def __init__(self, app):
        self.app = app

    def flush_routes(self):
        self._rtable = _routes.prepare(self.app.conf.CELERY_ROUTES)

    def Queues(self, queues, create_missing=None):
        """Create new :class:`Queues` instance, using queue defaults
        from the current configuration."""
        conf = self.app.conf
        if create_missing is None:
            create_missing = conf.CELERY_CREATE_MISSING_QUEUES
        if not queues and conf.CELERY_DEFAULT_QUEUE:
            queues = (Queue(conf.CELERY_DEFAULT_QUEUE,
                            exchange=self.default_exchange,
                            routing_key=conf.CELERY_DEFAULT_ROUTING_KEY), )
        return Queues(queues, self.default_exchange, create_missing)

    def Router(self, queues=None, create_missing=None):
        """Returns the current task router."""
        return _routes.Router(self.routes, queues or self.queues,
                              self.app.either('CELERY_CREATE_MISSING_QUEUES',
                                              create_missing), app=self.app)

    @cached_property
    def TaskConsumer(self):
        """Return consumer configured to consume from the queues
        we are configured for (``app.amqp.queues.consume_from``)."""
        return self.app.subclass_with_self(TaskConsumer,
                                           reverse='amqp.TaskConsumer')
    get_task_consumer = TaskConsumer  # XXX compat

    @cached_property
    def TaskProducer(self):
        """Returns publisher used to send tasks.

        You should use `app.send_task` instead.

        """
        conf = self.app.conf
        return self.app.subclass_with_self(TaskProducer,
                reverse='amqp.TaskProducer',
                exchange=self.default_exchange,
                routing_key=conf.CELERY_DEFAULT_ROUTING_KEY,
                serializer=conf.CELERY_TASK_SERIALIZER,
                compression=conf.CELERY_MESSAGE_COMPRESSION,
                retry=conf.CELERY_TASK_PUBLISH_RETRY,
                retry_policy=conf.CELERY_TASK_PUBLISH_RETRY_POLICY,
                utc=conf.CELERY_ENABLE_UTC)
    TaskPublisher = TaskProducer  # compat

    @cached_property
    def default_queue(self):
        return self.queues[self.app.conf.CELERY_DEFAULT_QUEUE]

    @cached_property
    def queues(self):
        """Queue name⇒ declaration mapping."""
        return self.Queues(self.app.conf.CELERY_QUEUES)

    @queues.setter  # noqa
    def queues(self, queues):
        return self.Queues(queues)

    @property
    def routes(self):
        if self._rtable is None:
            self.flush_routes()
        return self._rtable

    @cached_property
    def router(self):
        return self.Router()

    @cached_property
    def producer_pool(self):
        return ProducerPool(self.app.pool, limit=self.app.pool.limit,
                            Producer=self.TaskProducer)
    publisher_pool = producer_pool  # compat alias

    @cached_property
    def default_exchange(self):
        return Exchange(self.app.conf.CELERY_DEFAULT_EXCHANGE,
                        self.app.conf.CELERY_DEFAULT_EXCHANGE_TYPE)