response save on storage logic added WIP

config.yaml

@@ -1,5 +1,9 @@
 service:
   logging_level: $LOGGING_LEVEL_ROOT|DEBUG  # Logging level for service logger
+  response:
+    save: True  # file-to-storage upload
+    extension: ".NER_ENTITIES.json.gz"
+    # extension: ".OBJECTS.json.gz"
 
 probing_webserver:
   host: $PROBING_WEBSERVER_HOST|"0.0.0.0"  # Probe webserver address

pyinfra/callback.py

@@ -1,9 +1,13 @@
 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):
@@ -56,8 +60,15 @@ def make_callback_for_output_queue(json_wrapped_body_processor, output_queue_nam
 
     def callback(channel, method, _, body):
 
-        result = json_wrapped_body_processor(body)
-        channel.basic_publish(exchange="", routing_key=output_queue_name, body=result)
+        dossier_id, file_id, result = json_wrapped_body_processor(body)
+
+        if not CONFIG.service.response.save:
+            channel.basic_publish(exchange="", routing_key=output_queue_name, body=result)
+        else:
+            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

pyinfra/core.py

@@ -1,4 +1,5 @@
 import gzip
+import json
 import logging
 from operator import itemgetter
 
@@ -52,13 +53,14 @@ def make_analyzer(analysis_endpoint):
 
 
 def make_payload_processor(load_data, analyze_file):
-    @json_wrap
     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 predictions
+            return dossier_id, file_id, json.dumps(predictions)
         except (DataLoadingFailure, AnalysisFailure) as err:
             logging.warning(f"Processing of {payload} failed.")
             raise ProcessingFailure() from err

pyinfra/storage/storages.py

@@ -1,3 +1,4 @@
+from pyinfra.exceptions import UnknownStorageBackend
 from pyinfra.storage.adapters.azure import AzureStorageAdapter
 from pyinfra.storage.adapters.s3 import S3StorageAdapter
 from pyinfra.storage.clients.azure import get_azure_client
@@ -11,3 +12,14 @@ def get_azure_storage(config=None):
 
 def get_s3_storage(config=None):
     return Storage(S3StorageAdapter(get_s3_client(config)))
+
+def get_storage(storage_backend):
+
+    if storage_backend == "s3":
+        storage = get_s3_storage()
+    elif storage_backend == "azure":
+        storage = get_azure_storage()
+    else:
+        raise UnknownStorageBackend(f"Unknown storage backend '{storage_backend}'.")
+
+    return storage

pyinfra/utils/file.py

@@ -1,6 +1,12 @@
 """Defines utilities for different operations on files."""
 
+import gzip
+from io import BytesIO, StringIO
+import json
 import os
+from pathlib import Path
+import sys
+import tempfile
 
 from pyinfra.config import CONFIG
 
@@ -15,3 +21,9 @@ def dossier_id_and_file_id_to_compressed_storage_pdf_object_name(dossier_id, fil
     path_no_ext = os.path.join(dossier_id, file_id)
     pdf_object_name = produce_compressed_storage_pdf_object_name(path_no_ext)
     return pdf_object_name
+
+
+def upload_compressed_response(storage, bucket_name, dossier_id, file_id, result) -> None:
+    data = gzip.compress(result.encode())
+    path_gz = f"{dossier_id}/{file_id}{CONFIG.service.response.extension}"
+    storage.put_object(bucket_name, path_gz, data)

src/serve.py

@@ -11,22 +11,8 @@ from pyinfra.callback import (
 from pyinfra.config import CONFIG
 from pyinfra.consume import consume, ConsumerError
 from pyinfra.core import make_payload_processor, make_storage_data_loader, make_analyzer
-from pyinfra.exceptions import UnknownStorageBackend
 from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
-from pyinfra.storage.storages import get_azure_storage, get_s3_storage
-
-
-def get_storage():
-
-    storage_backend = CONFIG.storage.backend
-    if storage_backend == "s3":
-        storage = get_s3_storage()
-    elif storage_backend == "azure":
-        storage = get_azure_storage()
-    else:
-        raise UnknownStorageBackend(f"Unknown storage backend '{storage_backend}'.")
-
-    return storage
+from pyinfra.storage.storages import get_storage
 
 
 def republish(channel, body, n_current_attempts):
@@ -40,7 +26,7 @@ def republish(channel, body, n_current_attempts):
 
 def make_callback():
 
-    load_data = make_storage_data_loader(get_storage(), CONFIG.storage.bucket)
+    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)