From b4b060ac0e11a7b4f2e9a839300757e6ccb26539 Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Fri, 11 Mar 2022 11:43:27 +0100 Subject: [PATCH] added compression / decompression; jsons put on storage rather than only result data --- pyinfra/test/unit_tests/queue_visitor_test.py | 11 ++++---- pyinfra/visitor.py | 27 ++++++++++++++----- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/pyinfra/test/unit_tests/queue_visitor_test.py b/pyinfra/test/unit_tests/queue_visitor_test.py index 226e993..8a7e107 100644 --- a/pyinfra/test/unit_tests/queue_visitor_test.py +++ b/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" diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py index defe94d..69f75e0 100644 --- a/pyinfra/visitor.py +++ b/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