diff --git a/pyinfra/pyampq/queue_manager/pika_queue_manager.py b/pyinfra/pyampq/queue_manager/pika_queue_manager.py index 274caef..da56823 100644 --- a/pyinfra/pyampq/queue_manager/pika_queue_manager.py +++ b/pyinfra/pyampq/queue_manager/pika_queue_manager.py @@ -4,6 +4,7 @@ import logging import pika from pyinfra.config import CONFIG +from pyinfra.exceptions import ProcessingFailure from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle logger = logging.getLogger("pika") @@ -53,6 +54,14 @@ def get_connection(): 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): super().__init__(input_queue, output_queue) @@ -61,19 +70,43 @@ class PikaQueueManager(QueueManager): self.channel.queue_declare(input_queue) self.channel.queue_declare(output_queue) + 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)) + self.channel.basic_publish("", self._input_queue, json.dumps(request).encode()) - def publish_response(self, message, callback): + 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 - response = json.dumps(callback(json.loads(body))) - self.channel.basic_publish("", self._output_queue, response) - self.channel.basic_ack(frame.delivery_tag) + 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) diff --git a/pyinfra/test/unit_tests/consumer_test.py b/pyinfra/test/unit_tests/consumer_test.py index 8e4bf2d..f30996b 100644 --- a/pyinfra/test/unit_tests/consumer_test.py +++ b/pyinfra/test/unit_tests/consumer_test.py @@ -1,8 +1,10 @@ import gzip +import logging from operator import itemgetter import pytest +from pyinfra.exceptions import ProcessingFailure from pyinfra.pyampq.consumer import Consumer from pyinfra.visitor import get_object_descriptor, ForwardingStrategy @@ -17,13 +19,30 @@ def access_callback(): return itemgetter("fileId") -@pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session") +@pytest.fixture() +def items(): + + def inner(): + for i in range(3): + body = { + "dossierId": "folder", + "fileId": f"file{i}", + "targetFileExtension": "in.gz", + "responseFileExtension": "out.gz", + } + yield f"{i}".encode(), body + + return list(inner()) + + class TestConsumer: + @pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session") def test_consuming_empty_input_queue_does_not_put_anything_on_output_queue(self, consumer, queue_manager): queue_manager.clear() consumer.consume() assert queue_manager.output_queue.empty() + @pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session") def test_consuming_nonempty_input_queue_puts_messages_on_output_queue_in_fifo_order( self, consumer, queue_manager, callback ): @@ -42,33 +61,55 @@ class TestConsumer: assert queue_manager.output_queue.to_list() == ["00", "11", "22"] + @pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session") @pytest.mark.parametrize("client_name", ["mock", "s3", "azure"], scope="session") @pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session") def test_consuming_nonempty_input_queue_with_visitor_puts_messages_on_output_queue_in_fifo_order( - self, consumer, queue_manager, visitor, bucket_name, storage + self, consumer, queue_manager, visitor, bucket_name, storage, items ): - def produce_items(): - for i in range(3): - body = { - "dossierId": "folder", - "fileId": f"file{i}", - "targetFileExtension": "in.gz", - "responseFileExtension": "out.gz", - } - yield f"{i}".encode(), body visitor.response_strategy = ForwardingStrategy() queue_manager.clear() storage.clear_bucket(bucket_name) - for data, message in produce_items(): + for data, message in items: storage.put_object(**get_object_descriptor(message), data=gzip.compress(data)) queue_manager.publish_request(message) requests = consumer.consume() - for _, r in zip(produce_items(), requests): + for _, r in zip(items, requests): queue_manager.publish_response(r, visitor) assert list(map(itemgetter("data"), queue_manager.output_queue.to_list())) == ["00", "11", "22"] + + @pytest.mark.parametrize("queue_manager_name", ["pika"], scope="session") + def test_message_is_republished_when_callback_raises_processing_failure_exception( + self, consumer, queue_manager, bucket_name, items + ): + + class DebugError(Exception): + pass + + def callback(_): + raise ProcessingFailure() + + def reject_patch(*args, **kwargs): + raise DebugError() + + queue_manager.reject = reject_patch + + queue_manager.clear() + + for data, message in items: + queue_manager.publish_request(message) + + requests = consumer.consume() + + logger = logging.getLogger("pyinfra.pyampq.queue_manager.pika_queue_manager") + logger.addFilter(lambda record: False) + + with pytest.raises(DebugError): + while True: + queue_manager.publish_response(next(requests), callback)