Marmelator/pyinfra
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 36 Wiki Activity

added republishing logic

Browse Source
This commit is contained in:
Matthias Bisping 2022-03-14 20:22:57 +01:00
parent a40f07139a
commit f8b60aad16
2 changed files with 92 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
pyinfra/pyampq/queue_manager/pika_queue_manager.py
View File

@ -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)

67
pyinfra/test/unit_tests/consumer_test.py
View File

@ -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)