added compression / decompression; jsons put on storage rather than only result data

pyinfra/test/unit_tests/queue_visitor_test.py

@@ -1,3 +1,4 @@
+import gzip
 import json
 
 import pytest
@@ -12,7 +13,7 @@ def visitor(storage, analysis_callback, response_strategy):
 
 @pytest.fixture()
 def body():
-    return {"dossierId": "folder", "fileId": "file", "targetFileExtension": "in", "responseFileExtension": "out"}
+    return {"dossierId": "folder", "fileId": "file", "targetFileExtension": "in.gz", "responseFileExtension": "out.gz"}
 
 
 @pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
@@ -21,13 +22,13 @@ class TestVisitor:
 
     def test_given_a_input_queue_message_callback_pulls_the_data_from_storage(self, visitor, body, storage, bucket_name):
         storage.clear_bucket(bucket_name)
-        storage.put_object(**get_object_descriptor(body), data=b"content")
+        storage.put_object(**get_object_descriptor(body), data=gzip.compress(b"content"))
         data_received = visitor.load_data(body)
         assert b"content" == data_received
 
     def test_visitor_pulls_and_processes_data(self, visitor, body, storage, bucket_name):
         storage.clear_bucket(bucket_name)
-        storage.put_object(**get_object_descriptor(body), data="2".encode())
+        storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode()))
         response_body = visitor.load_and_process(body)
         assert response_body["data"] == "22"
 
@@ -37,7 +38,7 @@ class TestVisitor:
         visitor.response_strategy.storage = storage
 
         storage.clear_bucket(bucket_name)
-        storage.put_object(**get_object_descriptor(body), data="2".encode())
+        storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode()))
         response_body = visitor(body)
         assert "data" not in response_body
-        assert json.loads(storage.get_object(**get_response_object_descriptor(body))) == "22"
+        assert json.loads(gzip.decompress(storage.get_object(**get_response_object_descriptor(body))))["data"] == "22"

pyinfra/visitor.py

@@ -1,9 +1,12 @@
 import abc
+import gzip
 import json
+import logging
 from operator import itemgetter
 from typing import Callable
 
 from pyinfra.config import CONFIG
+from pyinfra.exceptions import DataLoadingFailure
 from pyinfra.storage.storage import Storage
 
 
@@ -44,7 +47,10 @@ class ResponseStrategy(abc.ABC):
 
 class StorageStrategy(ResponseStrategy):
     def handle_response(self, body):
-        self.storage.put_object(**get_response_object_descriptor(body), data=json.dumps(body.pop("data")).encode())
+        self.storage.put_object(
+            **get_response_object_descriptor(body), data=gzip.compress(json.dumps(body).encode())
+        )
+        body.pop("data")
         return body
 
 
@@ -63,7 +69,20 @@ class QueueVisitor:
         self.response_strategy.storage = self.storage
 
     def load_data(self, body):
-        return self.storage.get_object(**get_object_descriptor(body))
+
+        def download():
+            logging.debug(f"Downloading {object_descriptor}...")
+            data = self.storage.get_object(**object_descriptor)
+            logging.debug(f"Downloaded {object_descriptor}.")
+            return data
+
+        object_descriptor = get_object_descriptor(body)
+
+        try:
+            return gzip.decompress(download())
+        except Exception as err:
+            logging.warning(f"Loading data from storage failed for {object_descriptor}.")
+            raise DataLoadingFailure() from err
 
     def process_data(self, data):
         return self.callback(data)
@@ -76,7 +95,3 @@ class QueueVisitor:
     def __call__(self, body):
         result_body = self.load_and_process(body)
         return self.response_strategy(result_body)
-
-
-# TODO: 1) add compression / decompression to visitor
-#       2) correct keys and values for storage response