Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
Toggle table of contents sidebar

# Source code for aio_pika.patterns.master

    import asyncio
    import gzip
    import json
    import logging
    from functools import partial
    from types import MappingProxyType
    from typing import Any, Awaitable, Mapping, Optional

    import aiormq

    from aio_pika.abc import (
        AbstractChannel, AbstractExchange, AbstractIncomingMessage, AbstractQueue,
        ConsumerTag, DeliveryMode,
    )
    from aio_pika.message import Message

    from ..tools import create_task, ensure_awaitable
    from .base import Base, CallbackType, Proxy, T

    log = logging.getLogger(__name__)

    class MessageProcessingError(Exception):
        pass

    class NackMessage(MessageProcessingError):
        def __init__(self, requeue: bool = False):
            self.requeue = requeue

    class RejectMessage(MessageProcessingError):
        def __init__(self, requeue: bool = False):
            self.requeue = requeue

    class Worker:
        __slots__ = (
            "queue",
            "consumer_tag",
            "loop",
        )

        def __init__(
            self, queue: AbstractQueue, consumer_tag: ConsumerTag,
            loop: asyncio.AbstractEventLoop,
        ):
            self.queue = queue
            self.consumer_tag = consumer_tag
            self.loop = loop

        def close(self) -> Awaitable[None]:
            """ Cancel subscription to the channel

            :return: :class:`asyncio.Task`
            """

            async def closer() -> None:
                await self.queue.cancel(self.consumer_tag)

            return create_task(closer)

    class Master(Base):
        __slots__ = (
            "channel",
            "loop",
            "proxy",
            "_requeue",
            "_reject_on_redelivered",
        )

        DELIVERY_MODE = DeliveryMode.PERSISTENT

        __doc__ = """
        Implements Master/Worker pattern.
        Usage example:

        `worker.py` ::

            master = Master(channel)
            worker = await master.create_worker('test_worker', lambda x: print(x))

        `master.py` ::

            master = Master(channel)
            await master.proxy.test_worker('foo')
        """

        def __init__(
            self,
            channel: AbstractChannel,
            requeue: bool = True,
            reject_on_redelivered: bool = False,
        ):
            """ Creates a new :class:`Master` instance.

            :param channel: Initialized instance of :class:`aio_pika.Channel`
            """
            self.channel: AbstractChannel = channel
            self.loop: asyncio.AbstractEventLoop = asyncio.get_event_loop()
            self.proxy = Proxy(self.create_task)

            self.channel.return_callbacks.add(self.on_message_returned)

            self._requeue = requeue
            self._reject_on_redelivered = reject_on_redelivered

        @property
        def exchange(self) -> AbstractExchange:
            return self.channel.default_exchange

        @staticmethod
        def on_message_returned(
            channel: Optional[AbstractChannel],
            message: AbstractIncomingMessage,
        ) -> None:
            log.warning(
                "Message returned. Probably destination queue does not exists: %r",
                message,
            )

        def serialize(self, data: Any) -> bytes:
            """ Serialize data to the bytes.
            Uses `pickle` by default.
            You should overlap this method when you want to change serializer

            :param data: Data which will be serialized
            :returns: bytes
            """
            return super().serialize(data)

        def deserialize(self, data: bytes) -> Any:
            """ Deserialize data from bytes.
            Uses `pickle` by default.
            You should overlap this method when you want to change serializer

            :param data: Data which will be deserialized
            :returns: :class:`Any`
            """
            return super().deserialize(data)

        @classmethod
        async def execute(
            cls, func: CallbackType, kwargs: Any,
        ) -> T:
            kwargs = kwargs or {}

            if not isinstance(kwargs, dict):
                logging.error("Bad kwargs %r received for the %r", kwargs, func)
                raise RejectMessage(requeue=False)

            return await func(**kwargs)

        async def on_message(
            self, func: CallbackType,
            message: AbstractIncomingMessage,
        ) -> None:
            async with message.process(
                requeue=self._requeue,
                reject_on_redelivered=self._reject_on_redelivered,
                ignore_processed=True,
            ):
                try:
                    await self.execute(func, self.deserialize(message.body))
                except RejectMessage as e:
                    await message.reject(requeue=e.requeue)
                except NackMessage as e:
                    await message.nack(requeue=e.requeue)

        async def create_queue(
            self, queue_name: str, **kwargs: Any,
        ) -> AbstractQueue:
            return await self.channel.declare_queue(queue_name, **kwargs)

        async def create_worker(
            self, queue_name: str,
            func: CallbackType,
            **kwargs: Any,
        ) -> Worker:
            """ Creates a new :class:`Worker` instance. """
            queue = await self.create_queue(queue_name, **kwargs)
            consumer_tag = await queue.consume(
                partial(self.on_message, ensure_awaitable(func)),
            )

            return Worker(queue, consumer_tag, self.loop)

        async def create_task(
            self, channel_name: str,
            kwargs: Mapping[str, Any] = MappingProxyType({}),
            **message_kwargs: Any,
        ) -> Optional[aiormq.abc.ConfirmationFrameType]:

            """ Creates a new task for the worker """
            message = Message(
                body=self.serialize(kwargs),
                content_type=self.CONTENT_TYPE,
                delivery_mode=self.DELIVERY_MODE,
                **message_kwargs,
            )

            return await self.exchange.publish(
                message, channel_name, mandatory=True,
            )

    class JsonMaster(Master):
        SERIALIZER = json
        CONTENT_TYPE = "application/json"

        def serialize(self, data: Any) -> bytes:
            return self.SERIALIZER.dumps(data, ensure_ascii=False).encode()

    class CompressedJsonMaster(Master):
        SERIALIZER = json
        CONTENT_TYPE = "application/json;compression=gzip"
        COMPRESS_LEVEL = 6

        def serialize(self, data: Any) -> bytes:
            return gzip.compress(
                self.SERIALIZER.dumps(data, ensure_ascii=False).encode(),
                compresslevel=self.COMPRESS_LEVEL,
            )

        def deserialize(self, data: bytes) -> Any:
            return self.SERIALIZER.loads(gzip.decompress(data))