added republishing logic

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)
+        n_attempts = get_n_previous_attempts(properties) + 1
-        self.channel.basic_ack(frame.delivery_tag)
+
+        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)

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)