From f0d63fead0a0d98a8e043855a7d4e98ef89342f8 Mon Sep 17 00:00:00 2001 From: N0str Date: Thu, 22 Nov 2018 20:04:11 +0300 Subject: [PATCH] Use aio_pika --- components/broker.py | 15 +++++++++------ lib/rabbitmq.py | 31 +++++++++++-------------------- 2 files changed, 20 insertions(+), 26 deletions(-) diff --git a/components/broker.py b/components/broker.py index e69f00d..adcf72f 100644 --- a/components/broker.py +++ b/components/broker.py @@ -1,5 +1,7 @@ import logging +from aio_pika import IncomingMessage + from .api import API from ..lib.rabbitmq import init_receiver, add_message_to_queue @@ -19,12 +21,13 @@ def __init__(self, core, event_loop, application_name, queue_name): self.queue_name = queue_name self.api = API(self, application_name) - async def callback(self, channel, body, envelope, properties): - try: - logging.debug(" [x] Received %r" % body) - await self.api.process(body.decode("utf-8")) - except Exception as e: - logging.error("Broker callback error: {}".format(e)) + async def callback(self, message: IncomingMessage): + with message.process(): + try: + logging.debug(" [x] Received %r" % message.body) + await self.api.process(message.body.decode("utf-8")) + except Exception as e: + logging.error("Broker callback error: {}".format(e)) async def send(self, message, host='localhost'): diff --git a/lib/rabbitmq.py b/lib/rabbitmq.py index 9bda7f6..4f4abdd 100644 --- a/lib/rabbitmq.py +++ b/lib/rabbitmq.py @@ -1,25 +1,16 @@ -import aioamqp - -# This file will be copied to every module by bash script during initial setup. +import aio_pika async def init_receiver(callback, queue_name, host='localhost'): - transport, protocol = await aioamqp.connect() - channel = await protocol.channel() - await channel.queue_declare(queue_name=queue_name, passive=False, auto_delete=True) - await channel.basic_consume(callback, queue_name=queue_name) + connection = await aio_pika.connect_robust(host=host) + channel = await connection.channel() + await channel.set_qos(prefetch_count=1) + queue = await channel.declare_queue(queue_name, durable=True) + await queue.consume(callback) async def add_message_to_queue(data, queue_name, host='localhost'): - transport, protocol = await aioamqp.connect() - channel = await protocol.channel() - - await channel.queue_declare(queue_name=queue_name, passive=False, auto_delete=True) - - await channel.basic_publish( - payload=data, - exchange_name='', - routing_key=queue_name - ) - - await protocol.close() - transport.close() \ No newline at end of file + connection = await aio_pika.connect_robust(host=host) + channel = await connection.channel() + message = aio_pika.Message(data.encode(), delivery_mode=aio_pika.DeliveryMode.PERSISTENT) + await channel.default_exchange.publish(message, routing_key=queue_name) + await connection.close()