Judgement Day - The Great Refactoring

pyinfra/callback.py

@@ -1,96 +0,0 @@
-import json
-import logging
-import tempfile
-from time import sleep
-from pyinfra.config import CONFIG
-
-from pyinfra.exceptions import AnalysisFailure, DataLoadingFailure
-from pyinfra.rabbitmq import make_connection, make_channel, declare_queue
-from pyinfra.storage.storages import get_storage
-from pyinfra.utils.file import upload_compressed_response
-
-
-def make_retry_callback(republish, max_attempts):
-    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):
-        return n_attempts < max_attempts
-
-    def callback(channel, method, properties, body):
-
-        n_attempts = get_n_previous_attempts(properties) + 1
-
-        logging.error(f"Message failed to process {n_attempts}/{max_attempts} times: {body}")
-
-        if attempts_remain(n_attempts):
-            republish(channel, body, n_attempts)
-            channel.basic_ack(delivery_tag=method.delivery_tag)
-
-        else:
-            logging.exception(f"Adding to dead letter queue: {body}")
-            channel.basic_reject(delivery_tag=method.delivery_tag, requeue=False)
-
-    return callback
-
-
-def wrap_callback_in_retry_logic(callback, retry_callback):
-    def wrapped_callback(channel, method, properties, body):
-        try:
-            callback(channel, method, properties, body)
-        except (AnalysisFailure, DataLoadingFailure):
-            sleep(5)
-            retry_callback(channel, method, properties, body)
-
-    return wrapped_callback
-
-
-def json_wrap(body_processor):
-    def inner(payload):
-        return json.dumps(body_processor(json.loads(payload)))
-
-    return inner
-
-
-def make_callback_for_output_queue(json_wrapped_body_processor, output_queue_name):
-
-    connection = make_connection()
-    channel = make_channel(connection)
-    declare_queue(channel, output_queue_name)
-
-    def callback(channel, method, _, body):
-        """
-        response is dossier_id, file_id and analysis result, if CONFIG.service.response.type == "file" the response only
-         contains file_id and dossier_id and the analysis result will be written in a json file
-        Args:
-            channel:
-            method:
-            _:
-            body:
-
-        Returns:
-
-        """
-
-        dossier_id, file_id, result = json_wrapped_body_processor(body)
-
-        result_key = CONFIG.service.response.key if CONFIG.service.response.key else "result"
-        result = json.dumps({"dossierId": dossier_id, "fileId": file_id, result_key: result})
-
-        if CONFIG.service.response.type == "file":
-            upload_compressed_response(
-                get_storage(CONFIG.storage.backend), CONFIG.storage.bucket, dossier_id, file_id, result
-            )
-            result = json.dumps({"dossierId": dossier_id, "fileId": file_id})
-
-        channel.basic_publish(exchange="", routing_key=output_queue_name, body=result)
-        channel.basic_ack(delivery_tag=method.delivery_tag)
-
-    return callback
-
-
-def make_retry_callback_for_output_queue(json_wrapped_body_processor, output_queue_name, retry_callback):
-    callback = make_callback_for_output_queue(json_wrapped_body_processor, output_queue_name)
-    callback = wrap_callback_in_retry_logic(callback, retry_callback)
-
-    return callback

pyinfra/consume.py

@@ -1,43 +0,0 @@
-import logging
-from typing import Callable
-
-import pika
-from retry import retry
-
-from pyinfra.exceptions import ProcessingFailure
-from pyinfra.rabbitmq import make_connection, make_channel, declare_queue
-
-
-class ConsumerError(Exception):
-    pass
-
-
-@retry(ConsumerError, tries=3, delay=5, jitter=(1, 3))
-def consume(queue_name: str, on_message_callback: Callable):
-
-    connection = make_connection()
-    channel = make_channel(connection)
-    declare_queue(channel, queue_name)
-    logging.info("Started infrastructure.")
-
-    while True:
-        try:
-            logging.info("Waiting for messages...")
-            channel.basic_consume(queue=queue_name, auto_ack=False, on_message_callback=on_message_callback)
-            channel.start_consuming()
-
-        except pika.exceptions.ConnectionClosedByBroker as err:
-            logging.critical(f"Caught a channel error: {err}, stopping.")
-            break
-
-        except pika.exceptions.AMQPChannelError as err:
-            logging.critical(f"Caught a channel error: {err}, stopping.")
-            break
-
-        except pika.exceptions.AMQPConnectionError as err:
-            logging.info("No AMPQ-connection found, retrying...")
-            logging.debug(err)
-            continue
-
-        except ProcessingFailure as err:
-            raise ConsumerError(f"Error while consuming {queue_name}.") from err

pyinfra/core.py

@@ -1,70 +0,0 @@
-import gzip
-import json
-import logging
-from operator import itemgetter
-
-import requests
-
-from pyinfra.config import CONFIG
-from pyinfra.exceptions import DataLoadingFailure, AnalysisFailure, ProcessingFailure
-from pyinfra.utils.file import combine_dossier_id_and_file_id_and_extension
-
-
-def make_storage_data_loader(storage, bucket_name):
-    def get_object_name(payload: dict) -> str:
-        dossier_id, file_id = itemgetter("dossierId", "fileId")(payload)
-        object_name = combine_dossier_id_and_file_id_and_extension(
-            dossier_id, file_id, CONFIG.storage.target_file_extension
-        )
-        return object_name
-
-    def download(payload):
-        object_name = get_object_name(payload)
-        logging.debug(f"Downloading {object_name}...")
-        data = storage.get_object(bucket_name, object_name)
-        logging.debug(f"Downloaded {object_name}.")
-        return data
-
-    def decompress(data):
-        return gzip.decompress(data)
-
-    def load_data(payload):
-        try:
-            return decompress(download(payload))
-        except Exception as err:
-            logging.warning(f"Loading data from storage failed for {payload}.")
-            raise DataLoadingFailure() from err
-
-    return load_data
-
-
-def make_analyzer(analysis_endpoint):
-    def analyze(data):
-        try:
-            logging.debug(f"Requesting analysis from {analysis_endpoint}...")
-            analysis_response = requests.post(analysis_endpoint, data=data)
-            analysis_response.raise_for_status()
-            analysis_response = analysis_response.json()
-            logging.debug(f"Received response.")
-            return analysis_response
-        except Exception as err:
-            logging.warning("Exception caught when calling analysis endpoint.")
-            raise AnalysisFailure() from err
-
-    return analyze
-
-
-def make_payload_processor(load_data, analyze_file):
-    def process(payload: dict):
-        logging.info(f"Processing {payload}...")
-        try:
-            payload = json.loads(payload)
-            dossier_id, file_id = itemgetter("dossierId", "fileId")(payload)
-            data = load_data(payload)
-            predictions = analyze_file(data)
-            return dossier_id, file_id, predictions
-        except (DataLoadingFailure, AnalysisFailure) as err:
-            logging.warning(f"Processing of {payload} failed.")
-            raise ProcessingFailure() from err
-
-    return process

pyinfra/exceptions.py

@@ -20,3 +20,7 @@ class InvalidEndpoint(ValueError):
 
 class UnknownClient(ValueError):
     pass
+
+
+class ConsumerError(Exception):
+    pass

pyinfra/pyampq/connection/connection.py

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

pyinfra/pyampq/connection/mock_connection.py

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

pyinfra/pyampq/connection/pika_connection.py

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

pyinfra/queue/consumer.py

@@ -1,4 +1,4 @@
-from pyinfra.pyampq.queue_manager.queue_manager import QueueManager
+from pyinfra.queue.queue_manager.queue_manager import QueueManager
 
 
 class Consumer:

pyinfra/queue/queue_manager/pika_queue_manager.py

@@ -5,7 +5,7 @@ import pika
 
 from pyinfra.config import CONFIG
 from pyinfra.exceptions import ProcessingFailure
-from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle
+from pyinfra.queue.queue_manager.queue_manager import QueueHandle, QueueManager
 
 logger = logging.getLogger("pika")
 logger.setLevel(logging.WARNING)

pyinfra/rabbitmq.py

@@ -1,31 +0,0 @@
-import pika
-
-from pyinfra.config import CONFIG
-
-
-def make_channel(connection) -> pika.adapters.blocking_connection.BlockingChannel:
-    channel = connection.channel()
-    channel.basic_qos(prefetch_count=CONFIG.rabbitmq.prefetch_count)
-    return channel
-
-
-def declare_queue(channel, queue: str):
-    args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": CONFIG.rabbitmq.queues.dead_letter}
-    return channel.queue_declare(queue=queue, auto_delete=False, arguments=args)
-
-
-def read_connection_params():
-    credentials = pika.PlainCredentials(CONFIG.rabbitmq.user, CONFIG.rabbitmq.password)
-    parameters = pika.ConnectionParameters(
-        host=CONFIG.rabbitmq.host,
-        port=CONFIG.rabbitmq.port,
-        heartbeat=CONFIG.rabbitmq.heartbeat,
-        credentials=credentials,
-    )
-    return parameters
-
-
-def make_connection() -> pika.BlockingConnection:
-    parameters = read_connection_params()
-    connection = pika.BlockingConnection(parameters)
-    return connection

pyinfra/test/queue/queue_manager_mock.py

@@ -1,5 +1,5 @@
-from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle
-from pyinfra.test.queue_mock import QueueMock
+from pyinfra.queue.queue_manager.queue_manager import QueueManager, QueueHandle
+from pyinfra.test.queue.queue_mock import QueueMock
 
 
 def monkey_patch_queue_handle(queue) -> QueueHandle:

pyinfra/test/unit_tests/conftest.py

@@ -6,15 +6,15 @@ 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.queue.queue_manager.pika_queue_manager import PikaQueueManager
+from pyinfra.queue.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.queue.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, QueueVisitor

pyinfra/test/unit_tests/consumer_test.py

@@ -6,7 +6,7 @@ from operator import itemgetter
 import pytest
 
 from pyinfra.exceptions import ProcessingFailure
-from pyinfra.pyampq.consumer import Consumer
+from pyinfra.queue.consumer import Consumer
 from pyinfra.visitor import get_object_descriptor, ForwardingStrategy
 
 
@@ -114,7 +114,7 @@ class TestConsumer:
 
         requests = consumer.consume()
 
-        logger = logging.getLogger("pyinfra.pyampq.queue_manager.pika_queue_manager")
+        logger = logging.getLogger("pyinfra.queue.queue_manager.pika_queue_manager")
         logger.addFilter(lambda record: False)
 
         with pytest.raises(DebugError):

pyinfra/test/unit_tests/core_test.py

@@ -1,32 +0,0 @@
-import json
-from unittest.mock import patch
-
-from pyinfra.core import make_analyzer, make_payload_processor
-from pyinfra.test.config import CONFIG
-
-
-@patch("requests.post")
-def test_analyse_returns_analysis(mock_post, storage_data, mock_response):
-    mock_post.return_value = mock_response
-
-    analyze = make_analyzer(CONFIG.mock_analysis_endpoint)
-    response = analyze(storage_data)
-
-    assert response == storage_data
-
-
-@patch("requests.post")
-def test_process_returns_dossier_id_file_id_predictions(
-    mock_post, mock_make_load_data, storage_data, mock_response, mock_payload
-):
-    mock_post.return_value = mock_response
-
-    analyze = make_analyzer(CONFIG.mock_analysis_endpoint)
-    mock_load_data = mock_make_load_data
-
-    process = make_payload_processor(mock_load_data, analyze)
-    dossier_id, file_id, predictions = process(mock_payload)
-
-    assert dossier_id == json.loads(mock_payload)["dossierId"]
-    assert file_id == json.loads(mock_payload)["fileId"]
-    assert predictions == storage_data

pyinfra/utils/file.py

@@ -1,16 +0,0 @@
-"""Defines utilities for different operations on files."""
-
-import gzip
-import os
-
-from pyinfra.config import CONFIG
-
-
-def combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, extension):
-    return f"{dossier_id}/{file_id}{extension}"
-
-
-def upload_compressed_response(storage, bucket_name, dossier_id, file_id, result) -> None:
-    data = gzip.compress(result.encode())
-    path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, CONFIG.service.response.extension)
-    storage.put_object(bucket_name, path_gz, data)

pyinfra/utils/retry.py

@@ -1,96 +0,0 @@
-import logging
-import time
-from functools import partial, wraps
-from math import exp
-from typing import Tuple, Type, Callable
-
-
-class NoAttemptsLeft(Exception):
-    pass
-
-
-class MaxTimeoutReached(Exception):
-    pass
-
-
-class _MethodDecoratorAdaptor(object):
-    def __init__(self, decorator, func):
-        self.decorator = decorator
-        self.func = func
-
-    def __call__(self, *args, **kwargs):
-        return self.decorator(self.func)(*args, **kwargs)
-
-    def __get__(self, obj, objtype):
-        return partial(self.__call__, obj)
-
-
-def auto_adapt_to_methods(decorator):
-    """Allows you to use the same decorator on methods and functions,
-    hiding the self argument from the decorator."""
-
-    def adapt(func):
-        return _MethodDecoratorAdaptor(decorator, func)
-
-    return adapt
-
-
-def max_attempts(
-    n_attempts: int = 5, exceptions: Tuple[Type[Exception]] = None, timeout: float = 0.1, max_timeout: float = 10
-) -> Callable:
-    """Function decorator that attempts to run the wrapped function a certain number of times. Timeouts increase
-    exponentially according to `Tₖ ≔ t eᵏ`, where `t` is the timeout factor `timeout` and `k` is the attempt number.
-    If `∑ᵢ Tᵢ > mₜ` at the `i-th` attempt, where `mₜ` is the maximum timeout, then the function raises
-    MaxTimeoutReached. If `k > mₐ`, where `mₐ` is the maximum number of attempts allowed, then the function
-    raises NoAttemptsLeft.
-
-    Args:
-        n_attempts: Number of times to attempt running the wrapped function.
-        exceptions: Exceptions to catch for a re-attempt.
-        timeout: Timeout factor in seconds.
-        max_timeout: Maximum allowed timeout.
-
-    Raises:
-        MaxTimeoutReached
-        NoAttemptsLeft
-
-    Returns:
-        Wrapped function.
-    """
-    if not exceptions:
-        exceptions = (Exception,)
-    assert isinstance(exceptions, tuple)
-
-    @auto_adapt_to_methods
-    def inner(func):
-        @wraps(func)
-        def inner(*args, **kwargs):
-            def run_attempt(attempt, timeout_aggr=0):
-                if attempt:
-                    try:
-                        return func(*args, **kwargs)
-                    except exceptions as err:
-
-                        attempt_num = n_attempts - attempt + 1
-                        next_timeout = timeout * exp(attempt_num - 1)  # start with timeout * e^0 = timeout
-
-                        logging.warn(f"{func.__name__} failed; attempt {attempt_num} of {n_attempts}")
-
-                        time_left = max(0, max_timeout - timeout_aggr)
-                        if time_left:
-                            sleep_for = min(next_timeout, time_left)
-                            time.sleep(sleep_for)
-                            return run_attempt(attempt - 1, timeout_aggr + sleep_for)
-                        else:
-                            logging.exception(err)
-                            raise MaxTimeoutReached(
-                                f"{func.__name__} reached maximum timeout ({max_timeout}) after {attempt_num} attempts."
-                            )
-                else:
-                    raise NoAttemptsLeft(f"{func.__name__} failed {n_attempts} times; all attempts expended.")
-
-            return run_attempt(n_attempts)
-
-        return inner
-
-    return inner

scripts/manage_minio.py

@@ -7,7 +7,6 @@ from tqdm import tqdm
 
 from pyinfra.config import CONFIG
 from pyinfra.storage.storages import get_s3_storage
-from pyinfra.utils.file import combine_dossier_id_and_file_id_and_extension
 
 
 def parse_args():
@@ -58,3 +57,13 @@ if __name__ == "__main__":
 
     elif args.command == "purge":
         storage.clear_bucket(bucket_name)
+
+
+def combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, extension):
+    return f"{dossier_id}/{file_id}{extension}"
+
+
+def upload_compressed_response(storage, bucket_name, dossier_id, file_id, result) -> None:
+    data = gzip.compress(result.encode())
+    path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, CONFIG.service.response.extension)
+    storage.put_object(bucket_name, path_gz, data)

scripts/mock_client.py

@@ -1,7 +1,8 @@
 import json
 
+import pika
+
 from pyinfra.config import CONFIG
-from pyinfra.rabbitmq import make_channel, declare_queue, make_connection
 from pyinfra.storage.storages import get_s3_storage
 
 
@@ -20,6 +21,23 @@ def build_message_bodies():
         ).encode()
 
 
+def make_channel(connection) -> pika.adapters.blocking_connection.BlockingChannel:
+    channel = connection.channel()
+    channel.basic_qos(prefetch_count=CONFIG.rabbitmq.prefetch_count)
+    return channel
+
+
+def declare_queue(channel, queue: str):
+    args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": CONFIG.rabbitmq.queues.dead_letter}
+    return channel.queue_declare(queue=queue, auto_delete=False, arguments=args)
+
+
+def make_connection() -> pika.BlockingConnection:
+    parameters = read_connection_params()
+    connection = pika.BlockingConnection(parameters)
+    return connection
+
+
 if __name__ == "__main__":
 
     connection = make_connection()
@@ -34,3 +52,14 @@ if __name__ == "__main__":
     for method_frame, _, body in channel.consume(queue=CONFIG.rabbitmq.queues.output):
         print(f"Received {json.loads(body)}")
         channel.basic_ack(method_frame.delivery_tag)
+
+
+def read_connection_params():
+    credentials = pika.PlainCredentials(CONFIG.rabbitmq.user, CONFIG.rabbitmq.password)
+    parameters = pika.ConnectionParameters(
+        host=CONFIG.rabbitmq.host,
+        port=CONFIG.rabbitmq.port,
+        heartbeat=CONFIG.rabbitmq.heartbeat,
+        credentials=credentials,
+    )
+    return parameters

src/serve.py

@@ -5,11 +5,10 @@ import requests
 from retry import retry
 
 from pyinfra.config import CONFIG, make_art
-from pyinfra.consume import ConsumerError
-from pyinfra.exceptions import AnalysisFailure
+from pyinfra.exceptions import AnalysisFailure, ConsumerError
 from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
-from pyinfra.pyampq.consumer import Consumer
-from pyinfra.pyampq.queue_manager.pika_queue_manager import PikaQueueManager
+from pyinfra.queue.consumer import Consumer
+from pyinfra.queue.queue_manager.pika_queue_manager import PikaQueueManager
 from pyinfra.storage.storages import get_storage
 from pyinfra.visitor import QueueVisitor, StorageStrategy