From 21592342d5fd1bdf155b4314205cb5cd08cf38ee Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Tue, 8 Mar 2022 13:37:40 +0100 Subject: [PATCH] pika queue manager works now in basic version --- config.yaml | 3 +- pyinfra/pyampq/__init__.py | 0 pyinfra/pyampq/connection/__init__.py | 0 pyinfra/pyampq/connection/connection.py | 8 ++ pyinfra/pyampq/connection/mock_connection.py | 7 ++ pyinfra/pyampq/connection/pika_connection.py | 4 + pyinfra/pyampq/consumer.py | 13 +++ pyinfra/pyampq/queue_manager/__init__.py | 0 .../queue_manager/pika_queue_manager.py | 107 ++++++++++++++++++ pyinfra/pyampq/queue_manager/queue_manager.py | 49 ++++++++ pyinfra/test/queue_manager_mock.py | 42 +++++++ pyinfra/test/queue_mock.py | 5 + pyinfra/test/unit_tests/consumer_test.py | 41 +++++++ 13 files changed, 278 insertions(+), 1 deletion(-) create mode 100644 pyinfra/pyampq/__init__.py create mode 100644 pyinfra/pyampq/connection/__init__.py create mode 100644 pyinfra/pyampq/connection/connection.py create mode 100644 pyinfra/pyampq/connection/mock_connection.py create mode 100644 pyinfra/pyampq/connection/pika_connection.py create mode 100644 pyinfra/pyampq/consumer.py create mode 100644 pyinfra/pyampq/queue_manager/__init__.py create mode 100644 pyinfra/pyampq/queue_manager/pika_queue_manager.py create mode 100644 pyinfra/pyampq/queue_manager/queue_manager.py create mode 100644 pyinfra/test/queue_manager_mock.py create mode 100644 pyinfra/test/queue_mock.py create mode 100644 pyinfra/test/unit_tests/consumer_test.py diff --git a/config.yaml b/config.yaml index ed2c651..c044f4a 100755 --- a/config.yaml +++ b/config.yaml @@ -1,5 +1,5 @@ service: - logging_level: $LOGGING_LEVEL_ROOT|DEBUG # Logging level for service logger + logging_level: $LOGGING_LEVEL_ROOT|INFO # Logging level for service logger response: type: $RESPONSE_TYPE|"file" # Whether the analysis response is stored as file on storage or sent as stream extension: $RESPONSE_FILE_EXTENSION|".IMAGE_INFO.json.gz" # {.IMAGE_INFO.json.gz | .NER_ENTITIES.json.gz} @@ -37,6 +37,7 @@ storage: backend: $STORAGE_BACKEND|s3 # The type of storage to use {s3, azure} bucket: $STORAGE_BUCKET|"pyinfra-test-bucket" # The bucket / container to pull files specified in queue requests from + # TODO: Caller should specify exact file name, including extension! target_file_extension: $TARGET_FILE_EXTENSION|".ORIGIN.pdf.gz" # {.TEXT.json.gz | .ORIGIN.pdf.gz} Defines type of file to pull from storage s3: diff --git a/pyinfra/pyampq/__init__.py b/pyinfra/pyampq/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pyinfra/pyampq/connection/__init__.py b/pyinfra/pyampq/connection/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pyinfra/pyampq/connection/connection.py b/pyinfra/pyampq/connection/connection.py new file mode 100644 index 0000000..096325e --- /dev/null +++ b/pyinfra/pyampq/connection/connection.py @@ -0,0 +1,8 @@ +import abc + + +class Connection(abc.ABC): + + @abc.abstractmethod + def establish(self): + pass \ No newline at end of file diff --git a/pyinfra/pyampq/connection/mock_connection.py b/pyinfra/pyampq/connection/mock_connection.py new file mode 100644 index 0000000..9b03c4c --- /dev/null +++ b/pyinfra/pyampq/connection/mock_connection.py @@ -0,0 +1,7 @@ +from pyampq.connection.connection import Connection + + +class MockConnection(Connection): + + def establish(self): + pass diff --git a/pyinfra/pyampq/connection/pika_connection.py b/pyinfra/pyampq/connection/pika_connection.py new file mode 100644 index 0000000..ab4413d --- /dev/null +++ b/pyinfra/pyampq/connection/pika_connection.py @@ -0,0 +1,4 @@ +class PikaConnection: + + def establish(self): + pass \ No newline at end of file diff --git a/pyinfra/pyampq/consumer.py b/pyinfra/pyampq/consumer.py new file mode 100644 index 0000000..d01f270 --- /dev/null +++ b/pyinfra/pyampq/consumer.py @@ -0,0 +1,13 @@ +from pyinfra.pyampq.queue_manager.queue_manager import QueueManager + + +class Consumer: + def __init__(self, callback, queue_manager: QueueManager): + self.queue_manager = queue_manager + self.callback = callback + + def consume_and_publish(self): + yield from self.queue_manager.consume_and_publish(self.callback) + + def consume(self): + yield from self.queue_manager.consume() diff --git a/pyinfra/pyampq/queue_manager/__init__.py b/pyinfra/pyampq/queue_manager/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pyinfra/pyampq/queue_manager/pika_queue_manager.py b/pyinfra/pyampq/queue_manager/pika_queue_manager.py new file mode 100644 index 0000000..9736e1e --- /dev/null +++ b/pyinfra/pyampq/queue_manager/pika_queue_manager.py @@ -0,0 +1,107 @@ +import json +import logging + +import pika + +from pyinfra.config import CONFIG +from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle + +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) + print(message) + + 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(): + + credentials = pika.PlainCredentials( + username=CONFIG.rabbitmq.user, + password=CONFIG.rabbitmq.password + ) + + kwargs = { + "host": CONFIG.rabbitmq.host, + "port": CONFIG.rabbitmq.port, + "credentials": credentials + } + + parameters = pika.ConnectionParameters(**kwargs) + + connection = pika.BlockingConnection(parameters=parameters) + + return connection + + +class PikaQueueManager(QueueManager): + + def __init__(self, input_queue, output_queue): + super().__init__(input_queue, output_queue) + connection = get_connection() + self.channel = connection.channel() + self.channel.queue_declare(input_queue) + self.channel.queue_declare(output_queue) + + def publish_request(self, request): + self.channel.basic_publish("", self._input_queue, json.dumps(request).encode()) + + def publish_response(self, message, callback): + + frame, properties, body = message + response = json.dumps(callback(json.loads(body))).encode() + + self.channel.basic_publish("", self._output_queue, response) + self.channel.basic_ack(frame.delivery_tag) + + def pull_request(self): + return self.channel.basic_get(self._input_queue) + + def consume(self): + return self.channel.consume(self._input_queue) + + def consume_and_publish(self, callback): + + logger.info(f"Consuming with callback {callback.__name__}") + + for message in self.consume(): + self.publish_response(message, callback) + + def clear(self): + self.channel.queue_purge(self._input_queue) + self.channel.queue_purge(self._output_queue) + + @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/pyampq/queue_manager/queue_manager.py b/pyinfra/pyampq/queue_manager/queue_manager.py new file mode 100644 index 0000000..50bf916 --- /dev/null +++ b/pyinfra/pyampq/queue_manager/queue_manager.py @@ -0,0 +1,49 @@ +import abc + + +class QueueHandle: + + def empty(self) -> bool: + raise NotImplemented() + + def to_list(self) -> list: + raise NotImplemented() + + +class QueueManager(abc.ABC): + + def __init__(self, input_queue, output_queue): + self._input_queue = input_queue + self._output_queue = output_queue + + @abc.abstractmethod + def publish_request(self, request): + pass + + @abc.abstractmethod + def publish_response(self, response, callback): + pass + + @abc.abstractmethod + def pull_request(self): + pass + + @abc.abstractmethod + def consume(self): + pass + + @abc.abstractmethod + def clear(self): + pass + + @abc.abstractmethod + def input_queue(self) -> QueueHandle: + pass + + @abc.abstractmethod + def output_queue(self) -> QueueHandle: + pass + + @abc.abstractmethod + def consume_and_publish(self, callback): + pass diff --git a/pyinfra/test/queue_manager_mock.py b/pyinfra/test/queue_manager_mock.py new file mode 100644 index 0000000..ba8a634 --- /dev/null +++ b/pyinfra/test/queue_manager_mock.py @@ -0,0 +1,42 @@ +from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle +from pyinfra.test.queue_mock import QueueMock + + +def monkey_patch_queue_handle(queue) -> QueueHandle: + queue_handle = QueueHandle() + queue_handle.empty = lambda: not queue + queue_handle.to_list = lambda: list(queue) + return queue_handle + + +class QueueManagerMock(QueueManager): + + def __init__(self, input_queue, output_queue): + super().__init__(QueueMock(), QueueMock()) + + def publish_request(self, request): + self._input_queue.append(request) + + def publish_response(self, response): + self._output_queue.append(response) + + def pull_request(self): + return self._input_queue.pop() + + def consume(self, callback): + while self._input_queue: + request = self._input_queue.popleft() + response = callback(request) + self._output_queue.append(response) + + def clear(self): + self._input_queue.clear() + self._output_queue.clear() + + @property + def input_queue(self) -> QueueHandle: + return monkey_patch_queue_handle(self._input_queue) + + @property + def output_queue(self) -> QueueHandle: + return monkey_patch_queue_handle(self._output_queue) diff --git a/pyinfra/test/queue_mock.py b/pyinfra/test/queue_mock.py new file mode 100644 index 0000000..c138f29 --- /dev/null +++ b/pyinfra/test/queue_mock.py @@ -0,0 +1,5 @@ +from collections import deque + + +class QueueMock(deque): + pass diff --git a/pyinfra/test/unit_tests/consumer_test.py b/pyinfra/test/unit_tests/consumer_test.py new file mode 100644 index 0000000..4215042 --- /dev/null +++ b/pyinfra/test/unit_tests/consumer_test.py @@ -0,0 +1,41 @@ +import pytest + +from pyinfra.pyampq.consumer import Consumer +from pyinfra.pyampq.queue_manager.pika_queue_manager import PikaQueueManager +from pyinfra.pyampq.queue_manager.queue_manager import QueueManager +from pyinfra.test.queue_manager_mock import QueueManagerMock + + +@pytest.fixture +def callback(): + return lambda x: x**2 + + +@pytest.fixture(scope="session") +def queue_manager() -> QueueManager: + # return QueueManagerMock("input", "output") + return PikaQueueManager("input", "output") + + +@pytest.fixture +def consumer(queue_manager, callback): + return Consumer(callback, queue_manager) + + +# def test_consuming_empty_input_queue_does_not_put_anything_on_output_queue(consumer, queue_manager): +# queue_manager.clear() +# consumer.consume() +# assert queue_manager.output_queue.empty() + +def test_consuming_nonempty_input_queue_puts_messages_on_output_queue_in_fifo_order(consumer, queue_manager, callback): + queue_manager.clear() + queue_manager.publish_request(1) + queue_manager.publish_request(2) + requests = consumer.consume() + + for i, r in enumerate(requests): + queue_manager.publish_response(r, callback) + if i == 1: + break + + assert queue_manager.output_queue.to_list() == [1, 4]