pika queue manager works now in basic version

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:

pyinfra/pyampq/connection/connection.py

@@ -0,0 +1,8 @@
+import abc
+
+
+class Connection(abc.ABC):
+
+    @abc.abstractmethod
+    def establish(self):
+        pass

pyinfra/pyampq/connection/mock_connection.py

@@ -0,0 +1,7 @@
+from pyampq.connection.connection import Connection
+
+
+class MockConnection(Connection):
+
+    def establish(self):
+        pass

pyinfra/pyampq/connection/pika_connection.py

@@ -0,0 +1,4 @@
+class PikaConnection:
+
+    def establish(self):
+        pass

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

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)

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

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)

pyinfra/test/queue_mock.py

@@ -0,0 +1,5 @@
+from collections import deque
+
+
+class QueueMock(deque):
+    pass

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]