added compression / decompression; jsons put on storage rather than only result data

Merge in RR/pyinfra from prometheus-tunneling to pika_encapsulation

Squashed commit of the following:

commit 350448ec0f14849844deaa1a86ba4397ab3ebf3c
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Wed Mar 9 16:52:02 2022 +0100

    quickfix

commit ee88be4c80abdf597013c743ce2d490ad2b3e029
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Wed Mar 9 15:26:19 2022 +0100

    added prometheus endpoint to tunnel metrics from analysis endpoint

config.yaml

@@ -1,9 +1,9 @@
 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}
+#    extension: $RESPONSE_FILE_EXTENSION|"IMAGE_INFO.json.gz"  # {.IMAGE_INFO.json.gz | .NER_ENTITIES.json.gz}
-    key: $RESPONSE_KEY|"imageMetadata" # the key of the result {result, imageMetadata}
+#    key: $RESPONSE_KEY|"imageMetadata" # the key of the result {result, imageMetadata}
 
 probing_webserver:
   host: $PROBING_WEBSERVER_HOST|"0.0.0.0"  # Probe webserver address
@@ -37,7 +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
-  target_file_extension: $TARGET_FILE_EXTENSION|".ORIGIN.pdf.gz"  # {.TEXT.json.gz | .ORIGIN.pdf.gz} Defines type of file to pull from storage
+  # TODO: Caller should specify exact file name, including extension!
 
   s3:
     endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000"

pyinfra/config.py

@@ -1,11 +1,23 @@
 """Implements a config object with dot-indexing syntax."""
 
-
 from envyaml import EnvYAML
 
 from pyinfra.locations import CONFIG_FILE
 
 
+def make_art():
+    return """
+______     _____       __           
+| ___ \   |_   _|     / _|          
+| |_/ /   _ | | _ __ | |_ _ __ __ _ 
+|  __/ | | || || '_ \|  _| '__/ _` |
+| |  | |_| || || | | | | | | | (_| |
+\_|   \__, \___/_| |_|_| |_|  \__,_|
+       __/ |                        
+      |___/                         
+"""
+
+
 def _get_item_and_maybe_make_dotindexable(container, item):
     ret = container[item]
     return DotIndexable(ret) if isinstance(ret, dict) else ret

pyinfra/flask.py

@@ -1,11 +1,12 @@
+import requests
 from flask import Flask, jsonify
+from retry import retry
 from waitress import serve
 
 from pyinfra.config import CONFIG
 
 
 def run_probing_webserver(app, host=None, port=None, mode=None):
-
     if not host:
         host = CONFIG.probing_webserver.host
 
@@ -23,7 +24,6 @@ def run_probing_webserver(app, host=None, port=None, mode=None):
 
 
 def set_up_probing_webserver():
-
     app = Flask(__name__)
 
     @app.route("/ready", methods=["GET"])
@@ -38,4 +38,14 @@ def set_up_probing_webserver():
         resp.status_code = 200
         return resp
 
+    @app.route("/prometheus", methods=["GET"])
+    def get_analysis_prometheus_endpoint():
+        @retry(requests.exceptions.ConnectionError, tries=3, delay=5, jitter=(1, 3))
+        def inner():
+            prom_endpoint = f"{CONFIG.rabbitmq.callback.analysis_endpoint}/prometheus"
+            metric = requests.get(prom_endpoint)
+            metric.raise_for_status()
+            return metric.text
+        return inner()
+
     return app

pyinfra/pyampq/connection/connection.py

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

pyinfra/pyampq/connection/mock_connection.py

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

pyinfra/pyampq/connection/pika_connection.py

@@ -0,0 +1,3 @@
+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,98 @@
+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)
+
+            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,47 @@
+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,43 @@
+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, message, callback):
+        self._output_queue.append(callback(message))
+
+    def pull_request(self):
+        return self._input_queue.popleft()
+
+    def consume(self):
+        while self._input_queue:
+            yield self.pull_request()
+
+    def consume_and_publish(self, callback):
+        for message in self.consume():
+            self.publish_response(message, callback)
+
+    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/conftest.py

@@ -6,20 +6,24 @@ import pytest
 
 from pyinfra.exceptions import UnknownClient
 from pyinfra.locations import TEST_DIR
+from pyinfra.pyampq.queue_manager.pika_queue_manager import PikaQueueManager
+from pyinfra.pyampq.queue_manager.queue_manager import QueueManager
 from pyinfra.storage.adapters.azure import AzureStorageAdapter
 from pyinfra.storage.adapters.s3 import S3StorageAdapter
 from pyinfra.storage.clients.azure import get_azure_client
 from pyinfra.storage.clients.s3 import get_s3_client
 from pyinfra.storage.storage import Storage
 from pyinfra.test.config import CONFIG
+from pyinfra.test.queue_manager_mock import QueueManagerMock
 from pyinfra.test.storage.adapter_mock import StorageAdapterMock
 from pyinfra.test.storage.client_mock import StorageClientMock
+from pyinfra.visitor import StorageStrategy, ForwardingStrategy
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def bucket_name():
     return "pyinfra-test-bucket"
 
@@ -51,7 +55,7 @@ def mock_make_load_data():
     return load_data
 
 
-@pytest.fixture(params=["minio", "aws"])
+@pytest.fixture(params=["minio", "aws"], scope="session")
 def storage(client_name, bucket_name, request):
     logger.debug("Setup for storage")
     storage = Storage(get_adapter(client_name, request.param))
@@ -71,3 +75,35 @@ def get_adapter(client_name, s3_backend):
         return S3StorageAdapter(get_s3_client(CONFIG.storage[s3_backend]))
     else:
         raise UnknownClient(client_name)
+
+
+def get_queue_manager(queue_manager_name) -> QueueManager:
+    if queue_manager_name == "mock":
+        return QueueManagerMock("input", "output")
+    if queue_manager_name == "pika":
+        return PikaQueueManager("input", "output")
+
+
+@pytest.fixture(scope="session")
+def queue_manager(queue_manager_name) -> QueueManager:
+    return get_queue_manager(queue_manager_name)
+
+
+@pytest.fixture
+def callback():
+    return lambda x: x * 2
+
+
+@pytest.fixture
+def analysis_callback(callback):
+    def inner(data: bytes):
+        return callback(data.decode())
+    return inner
+
+
+@pytest.fixture
+def response_strategy(response_strategy_name):
+    if response_strategy_name == "storage":
+        return StorageStrategy()
+    if response_strategy_name == "forwarding":
+        return ForwardingStrategy()

pyinfra/test/unit_tests/consumer_test.py

@@ -0,0 +1,36 @@
+import pytest
+
+from pyinfra.pyampq.consumer import Consumer
+
+
+@pytest.fixture
+def consumer(queue_manager, callback):
+    return Consumer(callback, queue_manager)
+
+
+@pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session")
+class TestConsumer:
+
+    def test_consuming_empty_input_queue_does_not_put_anything_on_output_queue(self, 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(
+            self, consumer, queue_manager, callback
+    ):
+
+        def produce_items():
+            return map(str, range(3))
+
+        queue_manager.clear()
+
+        for item in produce_items():
+            queue_manager.publish_request(item)
+
+        requests = consumer.consume()
+
+        for _, r in zip(produce_items(), requests):
+            queue_manager.publish_response(r, callback)
+
+        assert queue_manager.output_queue.to_list() == ["00", "11", "22"]

pyinfra/test/unit_tests/queue_visitor_test.py

@@ -0,0 +1,44 @@
+import gzip
+import json
+
+import pytest
+
+from pyinfra.visitor import QueueVisitor, get_object_descriptor, StorageStrategy, get_response_object_descriptor
+
+
+@pytest.fixture()
+def visitor(storage, analysis_callback, response_strategy):
+    return QueueVisitor(storage, analysis_callback, response_strategy)
+
+
+@pytest.fixture()
+def body():
+    return {"dossierId": "folder", "fileId": "file", "targetFileExtension": "in.gz", "responseFileExtension": "out.gz"}
+
+
+@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
+@pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session")
+class TestVisitor:
+
+    def test_given_a_input_queue_message_callback_pulls_the_data_from_storage(self, visitor, body, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
+        storage.put_object(**get_object_descriptor(body), data=gzip.compress(b"content"))
+        data_received = visitor.load_data(body)
+        assert b"content" == data_received
+
+    def test_visitor_pulls_and_processes_data(self, visitor, body, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
+        storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode()))
+        response_body = visitor.load_and_process(body)
+        assert response_body["data"] == "22"
+
+    def test_visitor_puts_response_on_storage(self, visitor, body, storage, bucket_name):
+
+        visitor.response_strategy = StorageStrategy()
+        visitor.response_strategy.storage = storage
+
+        storage.clear_bucket(bucket_name)
+        storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode()))
+        response_body = visitor(body)
+        assert "data" not in response_body
+        assert json.loads(gzip.decompress(storage.get_object(**get_response_object_descriptor(body))))["data"] == "22"

pyinfra/test/unit_tests/storage_test.py

@@ -9,7 +9,7 @@ logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
-@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"])
+@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
 class TestStorage:
     def test_clearing_bucket_yields_empty_bucket(self, storage, bucket_name):
         storage.clear_bucket(bucket_name)
@@ -17,26 +17,31 @@ class TestStorage:
         assert not {*data_received}
 
     def test_getting_object_put_in_bucket_is_object(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "file", b"content")
         data_received = storage.get_object(bucket_name, "file")
         assert b"content" == data_received
 
     def test_getting_nested_object_put_in_bucket_is_nested_object(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "folder/file", b"content")
         data_received = storage.get_object(bucket_name, "folder/file")
         assert b"content" == data_received
 
     def test_getting_objects_put_in_bucket_are_objects(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "file1", b"content 1")
         storage.put_object(bucket_name, "folder/file2", b"content 2")
         data_received = storage.get_all_objects(bucket_name)
         assert {b"content 1", b"content 2"} == {*data_received}
 
     def test_make_bucket_produces_bucket(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.make_bucket(bucket_name)
         assert storage.has_bucket(bucket_name)
 
     def test_listing_bucket_files_yields_all_files_in_bucket(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "file1", b"content 1")
         storage.put_object(bucket_name, "file2", b"content 2")
         full_names_received = storage.get_all_object_names(bucket_name)

pyinfra/visitor.py

@@ -0,0 +1,97 @@
+import abc
+import gzip
+import json
+import logging
+from operator import itemgetter
+from typing import Callable
+
+from pyinfra.config import CONFIG
+from pyinfra.exceptions import DataLoadingFailure
+from pyinfra.storage.storage import Storage
+
+
+def get_object_name(body):
+    dossier_id, file_id, target_file_extension = itemgetter("dossierId", "fileId", "targetFileExtension")(body)
+    object_name = f"{dossier_id}/{file_id}.{target_file_extension}"
+    return object_name
+
+
+def get_response_object_name(body):
+    dossier_id, file_id, response_file_extension = itemgetter("dossierId", "fileId", "responseFileExtension")(body)
+    object_name = f"{dossier_id}/{file_id}.{response_file_extension}"
+    return object_name
+
+
+def get_object_descriptor(body):
+    return {
+        "bucket_name": CONFIG.storage.bucket,
+        "object_name": get_object_name(body)
+    }
+
+
+def get_response_object_descriptor(body):
+    return {
+        "bucket_name": CONFIG.storage.bucket,
+        "object_name": get_response_object_name(body)
+    }
+
+
+class ResponseStrategy(abc.ABC):
+    @abc.abstractmethod
+    def handle_response(self, body):
+        pass
+
+    def __call__(self, body):
+        return self.handle_response(body)
+
+
+class StorageStrategy(ResponseStrategy):
+    def handle_response(self, body):
+        self.storage.put_object(
+            **get_response_object_descriptor(body), data=gzip.compress(json.dumps(body).encode())
+        )
+        body.pop("data")
+        return body
+
+
+class ForwardingStrategy(ResponseStrategy):
+    def handle_response(self, body):
+        return body
+
+
+class QueueVisitor:
+
+    def __init__(self, storage: Storage, callback: Callable, response_strategy):
+        self.storage = storage
+        self.callback = callback
+        self.response_strategy = response_strategy
+        # TODO: a little dirty
+        self.response_strategy.storage = self.storage
+
+    def load_data(self, body):
+
+        def download():
+            logging.debug(f"Downloading {object_descriptor}...")
+            data = self.storage.get_object(**object_descriptor)
+            logging.debug(f"Downloaded {object_descriptor}.")
+            return data
+
+        object_descriptor = get_object_descriptor(body)
+
+        try:
+            return gzip.decompress(download())
+        except Exception as err:
+            logging.warning(f"Loading data from storage failed for {object_descriptor}.")
+            raise DataLoadingFailure() from err
+
+    def process_data(self, data):
+        return self.callback(data)
+
+    def load_and_process(self, body):
+        data = self.process_data(self.load_data(body))
+        result_body = {**body, "data": data}
+        return result_body
+
+    def __call__(self, body):
+        result_body = self.load_and_process(body)
+        return self.response_strategy(result_body)

src/serve.py

@@ -8,7 +8,7 @@ from pyinfra.callback import (
     make_retry_callback,
     make_callback_for_output_queue,
 )
-from pyinfra.config import CONFIG
+from pyinfra.config import CONFIG, make_art
 from pyinfra.consume import consume, ConsumerError
 from pyinfra.core import make_payload_processor, make_storage_data_loader, make_analyzer
 from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
@@ -49,6 +49,7 @@ def make_callback():
 
 def main():
     # TODO: implement meaningful checks
+    logging.info(make_art())
     webserver = Process(target=run_probing_webserver, args=(set_up_probing_webserver(),))
     logging.info("Starting webserver...")
     webserver.start()