refactoring: added cached pipeline factory

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):
+    def __init__(self, pipeline_factory: CachedPipelineFactory):
-        self.base_url = base_url
         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:
+        return self.pipeline_factory.get_pipeline(endpoint)
-            pipeline = self.endpoint2pipeline[endpoint]
-
-        else:
-            pipeline = self.pipeline_factory(endpoint)
-            self.endpoint2pipeline[endpoint] = pipeline
-
-        return pipeline
 
     @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(operation)
-        pipeline = self.__get_pipeline(endpoint)
 
         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}'.")

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"])

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}"