added list bucket files & refactor serve.py to take new storage logic

config.yaml

@@ -37,7 +37,6 @@ probing_webserver:
   port: $PROBING_WEBSERVER_PORT|8080  # Probe webserver port
   mode: $PROBING_WEBSERVER_MODE|production  # webserver mode: {development, production}
 
-
 test:
   minio:
     s3:
@@ -54,3 +53,5 @@ test:
   azure:
     azure_blob_storage:
       connection_string: "DefaultEndpointsProtocol=https;AccountName=iqserdevelopment;AccountKey=4imAbV9PYXaztSOMpIyAClg88bAZCXuXMGJG0GA1eIBpdh2PlnFGoRBnKqLy2YZUSTmZ3wJfC7tzfHtuC6FEhQ==;EndpointSuffix=core.windows.net"
+
+  bucket: "pyinfra-test-bucket"

pyinfra/storage/adapters/adapter.py

@@ -28,3 +28,7 @@ class StorageAdapter(ABC):
     @abstractmethod
     def clear_bucket(self, bucket_name):
         pass
+
+    @abstractmethod
+    def list_bucket_files(self, bucket_name):
+        pass

pyinfra/storage/adapters/azure.py

@@ -1,5 +1,7 @@
+from itertools import repeat
 import logging
 from functools import wraps
+from operator import attrgetter
 
 from azure.storage.blob import ContainerClient, BlobServiceClient
 from retry import retry
@@ -70,3 +72,8 @@ class AzureStorageAdapter(StorageAdapter):
         container_client = self.__client.get_container_client(bucket_name)
         blobs = container_client.list_blobs()
         container_client.delete_blobs(*blobs)
+
+    def list_bucket_files(self, bucket_name):
+        container_client = self.__provide_container_client(bucket_name)
+        blobs = container_client.list_blobs()
+        return zip(repeat(bucket_name), map(attrgetter("name"), blobs))

pyinfra/storage/adapters/s3.py

@@ -1,5 +1,7 @@
 import io
+from itertools import repeat
 import logging
+from operator import attrgetter
 
 from minio import Minio
 
@@ -47,3 +49,7 @@ class S3StorageAdapter(StorageAdapter):
         objects = self.__client.list_objects(bucket_name)
         for obj in objects:
             self.__client.remove_object(bucket_name, obj.object_name)
+
+    def list_bucket_files(self, bucket_name):
+        objs = self.__client.list_objects(bucket_name)
+        return zip(repeat(bucket_name), map(attrgetter("object_name"), objs))

pyinfra/storage/storage.py

@@ -20,6 +20,9 @@ class Storage:
     def clear_bucket(self, bucket_name):
         return self.__adapter.clear_bucket(bucket_name)
 
+    def list_bucket_files(self, bucket_name):
+        return self.__adapter.list_bucket_files(bucket_name)
+
 
 # import abc
 # import gzip

pyinfra/test/storage/adapter_mock.py

@@ -25,3 +25,6 @@ class StorageAdapterMock(StorageAdapter):
 
     def clear_bucket(self, bucket_name):
         return self.__client.clear_bucket(bucket_name)
+
+    def list_bucket_files(self, bucket_name):
+        return self.__client.list_bucket_files(bucket_name)

pyinfra/test/storage/client_mock.py

@@ -1,3 +1,6 @@
+from itertools import repeat
+
+
 class StorageClientMock:
     def __init__(self):
         self.__data = {}
@@ -19,3 +22,6 @@ class StorageClientMock:
 
     def clear_bucket(self, bucket_name):
         self.__data[bucket_name] = {}
+
+    def list_bucket_files(self, bucket_name):
+        return zip(repeat(bucket_name), self.__data[bucket_name])

pyinfra/test/unit_tests/storage_test.py

@@ -36,6 +36,12 @@ class TestStorage:
         storage.make_bucket(bucket_name)
         assert storage.bucket_exists(bucket_name)
 
+    def test_listing_bucket_files_yields_all_files_in_bucket(self, storage, bucket_name):
+        storage.put(bucket_name, "file1", b"content 1")
+        storage.put(bucket_name, "file2", b"content 2")
+        full_names_recived = storage.list_bucket_files(bucket_name)
+        assert {(bucket_name, "file1"), (bucket_name, "file2")} == {*full_names_recived}
+
 
 def get_adapter(client_name, s3_backend):
     if client_name == "mock":

pytest.ini

@@ -1,5 +1,4 @@
 [pytest]
 log_cli = 1
 log_cli_level = DEBUG
-addopts = --docker-compose=docker-compose.yaml
 

scripts/mock_client.py

@@ -1,13 +1,13 @@
 import json
 
 from pyinfra.config import CONFIG
-from pyinfra.storage.minio import MinioHandle
 from pyinfra.rabbitmq import make_channel, declare_queue, make_connection
+from pyinfra.storage.storages import get_s3_storage
 
 
 def build_message_bodies():
-    minio_client = MinioHandle()
-    for dossier_id, pdf_name in minio_client.list_folders_and_files():
+    minio_client = get_s3_storage()
+    for dossier_id, pdf_name in minio_client.list_bucket_files(CONFIG.test.bucket):
         file_id = pdf_name.split(".")[0]
         yield json.dumps({"dossierId": dossier_id, "fileId": file_id})
 

src/serve.py

@@ -13,17 +13,16 @@ 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.azure_blob_storage import AzureBlobStorageHandle
-from pyinfra.storage.minio import MinioHandle
+from pyinfra.storage.storages import get_azure_storage, get_s3_storage
 
 
 def get_storage():
 
     storage_backend = CONFIG.service.storage_backend
     if storage_backend == "s3":
-        storage = MinioHandle()
+        storage = get_s3_storage()
     elif storage_backend == "azure":
-        storage = AzureBlobStorageHandle()
+        storage = get_azure_storage()
     else:
         raise UnknownStorageBackend(f"Unknown storage backend '{storage_backend}'.")