143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
import json
|
|
import logging
|
|
|
|
import pika
|
|
|
|
from pyinfra.config import CONFIG
|
|
from pyinfra.exceptions import ProcessingFailure
|
|
from pyinfra.queue.queue_manager.queue_manager import QueueHandle, QueueManager
|
|
|
|
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
|
|
|
|
|
|
def get_n_previous_attempts(props):
|
|
return 0 if props.headers is None else props.headers.get("x-retry-count", 0)
|
|
|
|
|
|
def attempts_remain(n_attempts, max_attempts):
|
|
return n_attempts < max_attempts
|
|
|
|
|
|
class PikaQueueManager(QueueManager):
|
|
def __init__(self, input_queue, output_queue, dead_letter_queue=None):
|
|
super().__init__(input_queue, output_queue)
|
|
|
|
self.connection = get_connection()
|
|
self.channel = self.connection.channel()
|
|
|
|
if not dead_letter_queue:
|
|
dead_letter_queue = CONFIG.rabbitmq.queues.dead_letter
|
|
|
|
args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": dead_letter_queue}
|
|
|
|
self.channel.queue_declare(input_queue, arguments=args, auto_delete=False)
|
|
self.channel.queue_declare(output_queue, arguments=args, auto_delete=False)
|
|
|
|
def republish(self, body, n_current_attempts, frame):
|
|
self.channel.basic_publish(
|
|
exchange="",
|
|
routing_key=self._input_queue,
|
|
body=body,
|
|
properties=pika.BasicProperties(headers={"x-retry-count": n_current_attempts}),
|
|
)
|
|
self.channel.basic_ack(delivery_tag=frame.delivery_tag)
|
|
|
|
def publish_request(self, request):
|
|
logger.debug(f"Publishing {request}")
|
|
self.channel.basic_publish("", self._input_queue, json.dumps(request).encode())
|
|
|
|
def reject(self, body, frame):
|
|
logger.exception(f"Adding to dead letter queue: {body}")
|
|
self.channel.basic_reject(delivery_tag=frame.delivery_tag, requeue=False)
|
|
|
|
def publish_response(self, message, callback, max_attempts=3):
|
|
|
|
logger.debug(f"Publishing response for {message}.")
|
|
|
|
frame, properties, body = message
|
|
|
|
n_attempts = get_n_previous_attempts(properties) + 1
|
|
|
|
try:
|
|
response = json.dumps(callback(json.loads(body)))
|
|
self.channel.basic_publish("", self._output_queue, response.encode())
|
|
self.channel.basic_ack(frame.delivery_tag)
|
|
except ProcessingFailure:
|
|
|
|
logger.error(f"Message failed to process {n_attempts}/{max_attempts} times: {body}")
|
|
|
|
if attempts_remain(n_attempts, max_attempts):
|
|
self.republish(body, n_attempts, frame)
|
|
else:
|
|
self.reject(body, frame)
|
|
|
|
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, visitor):
|
|
|
|
logger.info(f"Consuming with callback {visitor.callback.__name__}")
|
|
|
|
for message in self.consume():
|
|
self.publish_response(message, visitor)
|
|
|
|
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)
|