From 5e56917970962a2e69bbd66a324bdb4618c040bd Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Mon, 27 Jun 2022 12:52:36 +0200
Subject: [PATCH] minor refactoring

---
 pyinfra/callback.py      | 28 +++++++++++++++-------------
 pyinfra/server/server.py | 21 ++++++++++++---------
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/pyinfra/callback.py b/pyinfra/callback.py
index 0b3d817..d9ce53d 100644
--- a/pyinfra/callback.py
+++ b/pyinfra/callback.py
@@ -31,7 +31,7 @@ class Callback:
         return pipeline
 
     @staticmethod
-    def __run_pipeline(pipeline, body):
+    def __run_pipeline(pipeline, analysis_input: dict):
         """
         TODO: Since data and metadata are passed as singletons, there is no buffering and hence no batching happening
          within the pipeline. However, the queue acknowledgment logic needs to be changed in order to facilitate
@@ -43,19 +43,19 @@ class Callback:
          operates on singletons ([data], [metadata]).
         """
 
-        def combine_storage_item_metadata_with_queue_message_metadata(body):
-            return merge(body["metadata"], omit(body, ["data", "metadata"]))
+        def combine_storage_item_metadata_with_queue_message_metadata(analysis_input):
+            return merge(analysis_input["metadata"], omit(analysis_input, ["data", "metadata"]))
 
-        def remove_queue_message_metadata(result):
-            metadata = omit(result["metadata"], queue_message_keys(body))
-            return {**result, "metadata": metadata}
+        def remove_queue_message_metadata(analysis_result):
+            metadata = omit(analysis_result["metadata"], queue_message_keys(analysis_input))
+            return {**analysis_result, "metadata": metadata}
 
-        def queue_message_keys(body):
-            return {*body.keys()}.difference({"data", "metadata"})
+        def queue_message_keys(analysis_input):
+            return {*analysis_input.keys()}.difference({"data", "metadata"})
 
         try:
-            data = body["data"]
-            metadata = combine_storage_item_metadata_with_queue_message_metadata(body)
+            data = analysis_input["data"]
+            metadata = combine_storage_item_metadata_with_queue_message_metadata(analysis_input)
             analysis_response_stream = pipeline([data], [metadata])
             analysis_response_stream = lmap(remove_queue_message_metadata, analysis_response_stream)
             return analysis_response_stream
@@ -64,13 +64,15 @@ class Callback:
             logger.error(err)
             raise AnalysisFailure from err
 
-    def __call__(self, body: dict):
-        operation = body.get("operation", "submit")
+    def __call__(self, analysis_input: dict):
+        """data_metadata_pack: {'dossierId': ..., 'fileId': ..., 'pages': ..., 'operation': ...}"""
+        # TODO: hardcoding of "submit" is complectig callback with analysis server implementation detail
+        operation = analysis_input.get("operation", "submit")
         endpoint = self.__make_endpoint(operation)
         pipeline = self.__get_pipeline(endpoint)
 
         try:
             logging.debug(f"Requesting analysis from {endpoint}...")
-            return self.__run_pipeline(pipeline, body)
+            return self.__run_pipeline(pipeline, analysis_input)
         except AnalysisFailure:
             logging.warning(f"Exception caught when calling analysis endpoint {endpoint}.")
diff --git a/pyinfra/server/server.py b/pyinfra/server/server.py
index 7c7553a..c9a8484 100644
--- a/pyinfra/server/server.py
+++ b/pyinfra/server/server.py
@@ -1,4 +1,4 @@
-from functools import singledispatch
+from functools import singledispatch, lru_cache
 from typing import Dict, Callable, Union
 
 from flask import Flask, jsonify, request
@@ -51,17 +51,20 @@ def __stream_fn_to_processing_server(operation2stream_fn: dict, buffer_size):
 
 
 def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunction]):
-    app = Flask(__name__)
-    registry = CollectorRegistry(auto_describe=True)
-
-    operation2processor = {
-        op: LazyRestProcessor(fn, **build_endpoint_suffixes(op)) for op, fn in operation2function.items()
-    }
+    @lru_cache(maxsize=None)
+    def get_registry():
+        return CollectorRegistry(auto_describe=True)
 
     def make_summary_instance(op: str):
         op = op.replace("_pickup", "") if op != "pickup" else "default"
         service_name = CONFIG.service.name
-        return Summary(f"redactmanager_{service_name}_{op}_seconds", f"Time spent on {op}.", registry=registry)
+        return Summary(f"redactmanager_{service_name}_{op}_seconds", f"Time spent on {op}.", registry=get_registry())
+
+    app = Flask(__name__)
+
+    operation2processor = {
+        op: LazyRestProcessor(fn, **build_endpoint_suffixes(op)) for op, fn in operation2function.items()
+    }
 
     submit_operation2processor = {submit_suffix(op): prc for op, prc in operation2processor.items()}
     pickup_operation2processor = {pickup_suffix(op): prc for op, prc in operation2processor.items()}
@@ -83,7 +86,7 @@ def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunctio
 
     @app.route("/prometheus", methods=["GET"])
     def prometheus():
-        return generate_latest(registry=registry)
+        return generate_latest(registry=get_registry())
 
     @app.route("/<operation>", methods=["POST", "PATCH"])
     def submit(operation):