103 lines
2.9 KiB
Python
103 lines
2.9 KiB
Python
import json
|
|
import logging
|
|
|
|
import pika
|
|
|
|
from pyinfra.config import CONFIG
|
|
from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle
|
|
|
|
logger = logging.getLogger("pika")
|
|
logger.setLevel(logging.WARNING)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(CONFIG.service.logging_level)
|
|
|
|
|
|
def monkey_patch_queue_handle(channel, queue) -> QueueHandle:
|
|
|
|
empty_message = (None, None, None)
|
|
|
|
def is_empty_message(message):
|
|
return message == empty_message
|
|
|
|
queue_handle = QueueHandle()
|
|
queue_handle.empty = lambda: is_empty_message(channel.basic_get(queue))
|
|
|
|
def produce_items():
|
|
|
|
while True:
|
|
message = channel.basic_get(queue)
|
|
|
|
if is_empty_message(message):
|
|
break
|
|
|
|
method_frame, properties, body = message
|
|
channel.basic_ack(method_frame.delivery_tag)
|
|
yield json.loads(body)
|
|
|
|
queue_handle.to_list = lambda: list(produce_items())
|
|
|
|
return queue_handle
|
|
|
|
|
|
def get_connection():
|
|
|
|
credentials = pika.PlainCredentials(username=CONFIG.rabbitmq.user, password=CONFIG.rabbitmq.password)
|
|
|
|
kwargs = {"host": CONFIG.rabbitmq.host, "port": CONFIG.rabbitmq.port, "credentials": credentials}
|
|
|
|
parameters = pika.ConnectionParameters(**kwargs)
|
|
|
|
connection = pika.BlockingConnection(parameters=parameters)
|
|
|
|
return connection
|
|
|
|
|
|
class PikaQueueManager(QueueManager):
|
|
def __init__(self, input_queue, output_queue):
|
|
super().__init__(input_queue, output_queue)
|
|
self.connection = get_connection()
|
|
self.channel = self.connection.channel()
|
|
self.channel.queue_declare(input_queue)
|
|
self.channel.queue_declare(output_queue)
|
|
|
|
def publish_request(self, request):
|
|
logger.debug(f"Publishing {request}")
|
|
self.channel.basic_publish("", self._input_queue, json.dumps(request))
|
|
|
|
def publish_response(self, message, callback):
|
|
|
|
logger.debug(f"Publishing response for {message}.")
|
|
|
|
frame, properties, body = message
|
|
response = json.dumps(callback(json.loads(body)))
|
|
|
|
self.channel.basic_publish("", self._output_queue, response)
|
|
self.channel.basic_ack(frame.delivery_tag)
|
|
|
|
def pull_request(self):
|
|
return self.channel.basic_get(self._input_queue)
|
|
|
|
def consume(self):
|
|
logger.debug("Consuming")
|
|
return self.channel.consume(self._input_queue)
|
|
|
|
def consume_and_publish(self, callback):
|
|
|
|
logger.info(f"Consuming with callback {callback.__name__}")
|
|
|
|
for message in self.consume():
|
|
self.publish_response(message, callback)
|
|
|
|
def clear(self):
|
|
self.channel.queue_purge(self._input_queue)
|
|
self.channel.queue_purge(self._output_queue)
|
|
|
|
@property
|
|
def input_queue(self) -> QueueHandle:
|
|
return monkey_patch_queue_handle(self.channel, self._input_queue)
|
|
|
|
@property
|
|
def output_queue(self) -> QueueHandle:
|
|
return monkey_patch_queue_handle(self.channel, self._output_queue)
|