data loading/publishing adjusted to new storage logic

config.yaml

@@ -31,6 +31,8 @@ service:
   name: $SERVICE_NAME|pyinfra-service-v1  # Name of the service in the kubernetes cluster
   storage_backend: $STORAGE_BACKEND|s3  # The storage to pull files to be processed from
   analysis_endpoint: $ANALYSIS_ENDPOINT|"http://127.0.0.1:5000"
+  bucket_name: $STORAGE_BUCKET_NAME|"pyinfra-test-bucket"  # TODO Maybe refactor!
+  target_file_extension: $TARGET_FILE_EXTENSION|"pdf"  # Defines type of file pulled from storage 
 
 probing_webserver:
   host: $PROBING_WEBSERVER_HOST|"0.0.0.0"  # Probe webserver address

pyinfra/core.py

@@ -9,7 +9,7 @@ from pyinfra.exceptions import DataLoadingFailure, AnalysisFailure, ProcessingFa
 from pyinfra.utils.file import dossier_id_and_file_id_to_compressed_storage_pdf_object_name
 
 
-def make_storage_data_loader(storage):
+def make_storage_data_loader(storage, bucket_name):
     def get_object_name(payload: dict) -> str:
         dossier_id, file_id = itemgetter("dossierId", "fileId")(payload)
         object_name = dossier_id_and_file_id_to_compressed_storage_pdf_object_name(dossier_id, file_id)
@@ -18,7 +18,7 @@ def make_storage_data_loader(storage):
     def download(payload):
         object_name = get_object_name(payload)
         logging.debug(f"Downloading {object_name}...")
-        data = storage.get_object(object_name)
+        data = storage.get(bucket_name, object_name)
         logging.debug(f"Downloaded {object_name}.")
         return data
 

pyinfra/storage/adapters/s3.py

@@ -40,7 +40,7 @@ class S3StorageAdapter(StorageAdapter):
                 response.release_conn()
 
     def get_all(self, bucket_name):
-        for obj in self.__client.list_objects(bucket_name):
+        for obj in self.__client.list_objects(bucket_name, recursive=True):
             logger.debug(f"Downloading '{obj.object_name}'...")
             yield self.get(bucket_name, obj.object_name)
 
@@ -51,5 +51,5 @@ class S3StorageAdapter(StorageAdapter):
             self.__client.remove_object(bucket_name, obj.object_name)
 
     def list_bucket_files(self, bucket_name):
-        objs = self.__client.list_objects(bucket_name)
+        objs = self.__client.list_objects(bucket_name, recursive=True)
         return zip(repeat(bucket_name), map(attrgetter("object_name"), objs))

pyinfra/utils/file.py

@@ -2,8 +2,12 @@
 
 import os
 
+from pyinfra.config import CONFIG
 
-def produce_compressed_storage_pdf_object_name(path_no_ext, ext="pdf"):
+
+def produce_compressed_storage_pdf_object_name(path_no_ext, ext=None):
+    if not ext:
+        ext = CONFIG.service.target_file_extension
     return f"{path_no_ext}.ORIGIN.{ext}.gz"
 
 

scripts/manage_minio.py

@@ -8,6 +8,7 @@ from pyinfra.config import CONFIG
 
 from pyinfra.storage.storages import get_s3_storage
 from pyinfra.test.unit_tests.conftest import bucket_name
+from pyinfra.utils.file import dossier_id_and_file_id_to_compressed_storage_pdf_object_name
 
 
 def parse_args():
@@ -29,7 +30,7 @@ def parse_args():
 
 def add_file_compressed(storage, bucket_name, dossier_id, path) -> None:
 
-    path_gz = str(Path(dossier_id) / (Path(path).name + ".gz"))
+    path_gz = dossier_id_and_file_id_to_compressed_storage_pdf_object_name(dossier_id, Path(path).stem)
 
     with open(path, "rb") as f:
         data = gzip.compress(f.read())
@@ -40,6 +41,8 @@ if __name__ == "__main__":
 
     storage = get_s3_storage()
     bucket_name = CONFIG.test.bucket
+    if not storage.bucket_exists(bucket_name):
+        storage.make_bucket(bucket_name)
 
     args = parse_args()
 

scripts/mock_client.py

@@ -7,8 +7,10 @@ from pyinfra.storage.storages import get_s3_storage
 
 def build_message_bodies():
     storage = get_s3_storage()
-    for dossier_id, pdf_name in storage.list_bucket_files(CONFIG.test.bucket):
+    for bucket_name, pdf_name in storage.list_bucket_files(CONFIG.test.bucket):
+        print(pdf_name) 
         file_id = pdf_name.split(".")[0]
+        dossier_id, file_id = file_id.split("/")
         yield json.dumps({"dossierId": dossier_id, "fileId": file_id})
 
 

src/serve.py

@@ -40,7 +40,7 @@ def republish(channel, body, n_current_attempts):
 
 def make_callback():
 
-    load_data = make_storage_data_loader(get_storage())
+    load_data = make_storage_data_loader(get_storage(), CONFIG.service.bucket_name)
     analyze_file = make_analyzer(CONFIG.service.analysis_endpoint)
 
     json_wrapped_body_processor = make_payload_processor(load_data, analyze_file)