diff --git a/pyinfra/config.py b/pyinfra/config.py index 98eba40..4156d8b 100644 --- a/pyinfra/config.py +++ b/pyinfra/config.py @@ -52,3 +52,7 @@ class Config(object): # Connection string for Azure storage self.storage_azureconnectionstring = read_from_environment("STORAGE_AZURECONNECTIONSTRING", "DefaultEndpointsProtocol=...") + + +def get_config() -> Config: + return Config() diff --git a/pyinfra/queue/pika_queue_manager.py b/pyinfra/queue/pika_queue_manager.py deleted file mode 100644 index ad5191a..0000000 --- a/pyinfra/queue/pika_queue_manager.py +++ /dev/null @@ -1,168 +0,0 @@ -import json -import logging -import time - -import pika - -from pyinfra.config import CONFIG -from pyinfra.exceptions import ProcessingFailure, DataLoadingFailure -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_params(): - - credentials = pika.PlainCredentials(username=CONFIG.rabbitmq.user, password=CONFIG.rabbitmq.password) - kwargs = { - "host": CONFIG.rabbitmq.host, - "port": CONFIG.rabbitmq.port, - "credentials": credentials, - "heartbeat": CONFIG.rabbitmq.heartbeat, - } - parameters = pika.ConnectionParameters(**kwargs) - - return parameters - - -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, connection_params=None): - super().__init__(input_queue, output_queue) - - if not connection_params: - connection_params = get_connection_params() - - self.connection = pika.BlockingConnection(parameters=connection_params) - self.channel = self.connection.channel() - self.channel.basic_qos(prefetch_count=1) - - 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, durable=True) - self.channel.queue_declare(output_queue, arguments=args, auto_delete=False, durable=True) - - 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.error(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"Processing {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, DataLoadingFailure): - - 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, inactivity_timeout=None): - logger.debug("Consuming") - return self.channel.consume(self._input_queue, inactivity_timeout=inactivity_timeout) - - 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 basic_consume_and_publish(self, visitor): - - logger.info(f"Basic consuming with callback {visitor.callback.__name__}") - - def callback(channel, frame, properties, body): - message = (frame, properties, body) - return self.publish_response(message, visitor) - - consumer_tag = None - - try: - consumer_tag = self.channel.basic_consume(self._input_queue, callback) - self.channel.start_consuming() - finally: - if consumer_tag: - self.channel.basic_cancel(consumer_tag) - - def clear(self): - try: - self.channel.queue_purge(self._input_queue) - self.channel.queue_purge(self._output_queue) - except pika.exceptions.ChannelWrongStateError: - pass - - @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) diff --git a/pyinfra/queue/queue_manager.py b/pyinfra/queue/queue_manager.py new file mode 100644 index 0000000..e5db562 --- /dev/null +++ b/pyinfra/queue/queue_manager.py @@ -0,0 +1,94 @@ +import json +import logging + +import pika, pika.exceptions +from typing import Callable + +from pyinfra.config import get_config, Config + +CONFIG = get_config() + +logger = logging.getLogger("pika") +logger.setLevel(logging.WARNING) + +logger = logging.getLogger(__name__) +logger.setLevel(CONFIG.logging_level_root) + + +def get_connection_params(config: Config) -> pika.ConnectionParameters: + credentials = pika.PlainCredentials(username=config.rabbitmq_username, password=config.rabbitmq_password) + pika_connection_params = { + "host": config.rabbitmq_host, + "port": config.rabbitmq_port, + "credentials": credentials, + "heartbeat": config.rabbitmq_heartbeat, + } + + return pika.ConnectionParameters(**pika_connection_params) + + +def _get_n_previous_attempts(props): + return 0 if props.headers is None else props.headers.get("x-retry-count", 0) + + +class QueueManager(object): + def __init__(self, config: Config = CONFIG): + connection_params = get_connection_params(config) + + self._connection = pika.BlockingConnection(parameters=connection_params) + self._channel = self._connection.channel() + self._channel.basic_qos(prefetch_count=1) + + args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": CONFIG.dead_letter_queue} + + self._input_queue = config.request_queue + self._output_queue = config.response_queue + + self._channel.queue_declare(self._input_queue, arguments=args, auto_delete=False, durable=True) + self._channel.queue_declare(self._output_queue, arguments=args, auto_delete=False, durable=True) + + self._consumer_token = None + + def start_consuming(self, process_message_callback: Callable): + callback = self._create_queue_callback(process_message_callback) + + logger.info("Consuming from queue") + consumer_tag = None + try: + consumer_tag = self._channel.basic_consume(self._input_queue, callback) + logger.info(f"Registered with consumer-tag: {consumer_tag}") + self._channel.start_consuming() + finally: + logger.warning("An unhandled exception occurred while consuming messages. Consuming will stop.") + if consumer_tag: + logger.info(f"Cancelling subscription for consumer-tag: {consumer_tag}") + self._channel.basic_cancel(consumer_tag) + else: + logger.warning("No consumer-tag found, a possible subscription cannot be cancelled") + + def _create_queue_callback(self, process_message_callback: Callable): + def callback(_channel, frame, properties, body): + logger.info("Received message from queue") + logger.debug(f"Processing {(frame, properties, body)}.") + + try: + unpacked_message_body = json.loads(body) + + callback_result = process_message_callback(unpacked_message_body) + logger.info("Processed message, publishing result to result-queue") + + self._channel.basic_publish("", self._output_queue, json.dumps(callback_result).encode()) + + self._channel.basic_ack(frame.delivery_tag) + except Exception as ex: + n_attempts = _get_n_previous_attempts(properties) + 1 + logger.warning(f"Message failed to process, {n_attempts} attempts, error {str(ex)}: message: {body}") + + return callback + + def clear(self): + try: + self._channel.queue_purge(self._input_queue) + self._channel.queue_purge(self._output_queue) + except pika.exceptions.ChannelWrongStateError: + pass