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 import (
    AbstractChannel, AbstractExchange, AbstractIncomingMessage, AbstractQueue,
    ConsumerTag, DeliveryMode,
from aio_pika.message import Message, ReturnedMessage

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

log = logging.getLogger(__name__)

class MessageProcessingError(Exception):

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

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

[docs]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
[docs] 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)
[docs]class Master(Base): __slots__ = ( "channel", "loop", "proxy", ) DELIVERY_MODE = DeliveryMode.PERSISTENT __doc__ = """ Implements Master/Worker pattern. Usage example: `` :: master = Master(channel) worker = await master.create_worker('test_worker', lambda x: print(x)) `` :: 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` """ AbstractChannel = channel self.loop: asyncio.AbstractEventLoop = asyncio.get_event_loop() self.proxy = Proxy(self.create_task) self._requeue = requeue self._reject_on_redelivered = reject_on_redelivered @property def exchange(self) -> AbstractExchange: return @staticmethod def on_message_returned( channel: AbstractChannel, message: ReturnedMessage, ) -> None: log.warning( "Message returned. Probably destination queue does not exists: %r", message, )
[docs] 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)
[docs] 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, **kwargs)
[docs] 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)
[docs] async def create_task( self, channel_name: str, kwargs: Mapping[str, Any] = MappingProxyType({}), **message_kwargs: Any, ) -> Optional[]: """ 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 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))