From 3bf00d11cd041dff325b66f13fcd00d3ce96b8b5 Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Thu, 30 Jun 2022 12:47:57 +0200
Subject: [PATCH] refactoring: added cached pipeline factory

---
 pyinfra/callback.py          | 24 ++++++------------------
 pyinfra/component_factory.py |  3 ++-
 pyinfra/pipeline_factory.py  | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+), 19 deletions(-)
 create mode 100644 pyinfra/pipeline_factory.py

diff --git a/pyinfra/callback.py b/pyinfra/callback.py
index 754258c..cec2d75 100644
--- a/pyinfra/callback.py
+++ b/pyinfra/callback.py
@@ -3,6 +3,7 @@ import logging
 from funcy import merge, omit, lmap
 
 from pyinfra.exceptions import AnalysisFailure
+from pyinfra.pipeline_factory import CachedPipelineFactory
 
 logger = logging.getLogger(__name__)
 
@@ -12,23 +13,11 @@ class Callback:
     endpoint.
     """
 
-    def __init__(self, base_url, pipeline_factory):
-        self.base_url = base_url
+    def __init__(self, pipeline_factory: CachedPipelineFactory):
         self.pipeline_factory = pipeline_factory
-        self.endpoint2pipeline = {}
-
-    def __make_endpoint(self, operation):
-        return f"{self.base_url}/{operation}"
 
     def __get_pipeline(self, endpoint):
-        if endpoint in self.endpoint2pipeline:
-            pipeline = self.endpoint2pipeline[endpoint]
-
-        else:
-            pipeline = self.pipeline_factory(endpoint)
-            self.endpoint2pipeline[endpoint] = pipeline
-
-        return pipeline
+        return self.pipeline_factory.get_pipeline(endpoint)
 
     @staticmethod
     def __run_pipeline(pipeline, analysis_input: dict):
@@ -67,11 +56,10 @@ class Callback:
     def __call__(self, analysis_input: dict):
         """data_metadata_pack: {'dossierId': ..., 'fileId': ..., 'pages': ..., 'operation': ...}"""
         operation = analysis_input.get("operation", "")
-        endpoint = self.__make_endpoint(operation)
-        pipeline = self.__get_pipeline(endpoint)
+        pipeline = self.__get_pipeline(operation)
 
         try:
-            logging.debug(f"Requesting analysis from {endpoint}...")
+            logging.debug(f"Requesting analysis for operation '{operation}'...")
             return self.__run_pipeline(pipeline, analysis_input)
         except AnalysisFailure:
-            logging.warning(f"Exception caught when calling analysis endpoint {endpoint}.")
+            logging.warning(f"Exception caught when calling analysis endpoint for operation '{operation}'.")
diff --git a/pyinfra/component_factory.py b/pyinfra/component_factory.py
index c0b0e9d..564c554 100644
--- a/pyinfra/component_factory.py
+++ b/pyinfra/component_factory.py
@@ -6,6 +6,7 @@ from funcy import project, identity, rcompose
 from pyinfra.callback import Callback
 from pyinfra.config import parse_disjunction_string
 from pyinfra.file_descriptor_manager import FileDescriptorManager
+from pyinfra.pipeline_factory import CachedPipelineFactory
 from pyinfra.queue.consumer import Consumer
 from pyinfra.queue.queue_manager.pika_queue_manager import PikaQueueManager
 from pyinfra.server.client_pipeline import ClientPipeline
@@ -36,7 +37,7 @@ class ComponentFactory:
     def get_callback(self, analysis_base_url=None):
         analysis_base_url = analysis_base_url or self.config.rabbitmq.callback.analysis_endpoint
 
-        callback = Callback(analysis_base_url, self.get_pipeline)
+        callback = Callback(CachedPipelineFactory(base_url=analysis_base_url, pipeline_factory=self.get_pipeline))
 
         def wrapped(body):
             body_repr = project(body, ["dossierId", "fileId", "operation"])
diff --git a/pyinfra/pipeline_factory.py b/pyinfra/pipeline_factory.py
new file mode 100644
index 0000000..fce742e
--- /dev/null
+++ b/pyinfra/pipeline_factory.py
@@ -0,0 +1,18 @@
+class CachedPipelineFactory:
+    def __init__(self, base_url, pipeline_factory):
+        self.base_url = base_url
+        self.operation2pipeline = {}
+        self.pipeline_factory = pipeline_factory
+
+    def get_pipeline(self, operation: str):
+        pipeline = self.operation2pipeline.get(operation, None) or self.__register_pipeline(operation)
+        return pipeline
+
+    def __register_pipeline(self, operation):
+        endpoint = self.__make_endpoint(operation)
+        pipeline = self.pipeline_factory(endpoint)
+        self.operation2pipeline[operation] = pipeline
+        return pipeline
+
+    def __make_endpoint(self, operation):
+        return f"{self.base_url}/{operation}"