RED-4653: Removed all unused code from queue-manager
This commit is contained in:
Viktor Seifert
parent
15b23a0019
commit
33b26527dc
@ -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()
|
||||
|
||||
@ -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)
|
||||
94
pyinfra/queue/queue_manager.py
Normal file
94
pyinfra/queue/queue_manager.py
Normal file
@ -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
|
||||
Loading…
x
Reference in New Issue
Block a user