From 64d8e0bd15730898274c08d34f9c34fbac559422 Mon Sep 17 00:00:00 2001 From: Viktor Seifert Date: Mon, 25 Jul 2022 13:57:06 +0200 Subject: [PATCH] RED-4653: Moved queue cancellation to a separate method so that it can be called on application exit --- pyinfra/queue/queue_manager.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/pyinfra/queue/queue_manager.py b/pyinfra/queue/queue_manager.py index 9644fb1..f226e0b 100644 --- a/pyinfra/queue/queue_manager.py +++ b/pyinfra/queue/queue_manager.py @@ -53,18 +53,23 @@ class QueueManager(object): callback = self._create_queue_callback(process_message_callback) self.logger.info("Consuming from queue") - consumer_tag = None + self._consumer_token = None try: - consumer_tag = self._channel.basic_consume(self._input_queue, callback) - self.logger.info(f"Registered with consumer-tag: {consumer_tag}") + self._consumer_token = self._channel.basic_consume(self._input_queue, callback) + self.logger.info(f"Registered with consumer-tag: {self._consumer_token}") self._channel.start_consuming() finally: self.logger.warning("An unhandled exception occurred while consuming messages. Consuming will stop.") - if consumer_tag: - self.logger.info(f"Cancelling subscription for consumer-tag: {consumer_tag}") - self._channel.basic_cancel(consumer_tag) - else: - self.logger.warning("No consumer-tag found, a possible subscription cannot be cancelled") + self.cancel_queue_subscription() + + def cancel_queue_subscription(self): + if self._consumer_token: + self.logger.info(f"Cancelling subscription for consumer-tag: {self._consumer_token}") + self._channel.basic_cancel(self._consumer_token) + + self._consumer_token = None + else: + self.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):