From 0ce30cafbf15142e3f773d0e510b7d22927ee86d Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Tue, 15 Mar 2022 11:42:04 +0100
Subject: [PATCH] updated serving logic for new queue, visior and storage logic

---
 src/serve.py | 76 ++++++++++++++++++++++++++++------------------------
 1 file changed, 41 insertions(+), 35 deletions(-)

diff --git a/src/serve.py b/src/serve.py
index 3490f9e..c18dd5f 100644
--- a/src/serve.py
+++ b/src/serve.py
@@ -1,48 +1,39 @@
 import logging
 from multiprocessing import Process
 
-import pika
+import requests
+from retry import retry
 
-from pyinfra.callback import (
-    make_retry_callback_for_output_queue,
-    make_retry_callback,
-    make_callback_for_output_queue,
-)
 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.consume import ConsumerError
+from pyinfra.exceptions import AnalysisFailure
 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.storage.storages import get_storage
+from pyinfra.visitor import QueueVisitor, StorageStrategy
 
 
-def republish(channel, body, n_current_attempts):
-    channel.basic_publish(
-        exchange="",
-        routing_key=CONFIG.rabbitmq.queues.input,
-        body=body,
-        properties=pika.BasicProperties(headers={"x-retry-count": n_current_attempts}),
-    )
+def make_callback(analysis_endpoint):
+    def callback(message):
 
+        def perform_operation(operation):
+            endpoint = f"{analysis_endpoint}/{operation}"
+            try:
+                logging.debug(f"Requesting analysis from {endpoint}...")
+                analysis_response = requests.post(endpoint, data=message["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(f"Exception caught when calling analysis endpoint {endpoint}.")
+                raise AnalysisFailure() from err
 
-def make_callback():
-
-    load_data = make_storage_data_loader(get_storage(CONFIG.storage.backend), CONFIG.storage.bucket)
-    analyze_file = make_analyzer(CONFIG.rabbitmq.callback.analysis_endpoint)
-
-    json_wrapped_body_processor = make_payload_processor(load_data, analyze_file)
-
-    if CONFIG.rabbitmq.callback.retry.enabled:
-        retry_callback = make_retry_callback(republish, max_attempts=CONFIG.rabbitmq.callback.retry.max_attempts)
-
-        callback = make_retry_callback_for_output_queue(
-            json_wrapped_body_processor=json_wrapped_body_processor,
-            output_queue_name=CONFIG.rabbitmq.queues.output,
-            retry_callback=retry_callback,
-        )
-    else:
-        callback = make_callback_for_output_queue(
-            json_wrapped_body_processor=json_wrapped_body_processor, output_queue_name=CONFIG.rabbitmq.queues.output
-        )
+        operations = message.get("operations", ["/"])
+        results = map(perform_operation, operations)
+        result = dict(zip(operations, results))
+        return result
 
     return callback
 
@@ -54,8 +45,23 @@ def main():
     logging.info("Starting webserver...")
     webserver.start()
 
+    callback = make_callback(CONFIG.rabbitmq.callback.analysis_endpoint)
+    storage = get_storage(CONFIG.storage.backend)
+    response_strategy = StorageStrategy(storage)
+    visitor = QueueVisitor(storage, callback, response_strategy)
+
+    queue_manager = PikaQueueManager(CONFIG.rabbitmq.queues.input, CONFIG.rabbitmq.queues.output)
+
+    @retry(ConsumerError, tries=3, delay=5, jitter=(1, 3))
+    def consume():
+        try:
+            consumer = Consumer(visitor, queue_manager)
+            consumer.consume_and_publish()
+        except Exception as err:
+            raise ConsumerError from err
+
     try:
-        consume(CONFIG.rabbitmq.queues.input, make_callback())
+        consume()
     except KeyboardInterrupt:
         pass
     except ConsumerError: