diff --git a/config.yaml b/config.yaml index 6c21fad..75a0408 100755 --- a/config.yaml +++ b/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" diff --git a/pyinfra/storage/adapters/adapter.py b/pyinfra/storage/adapters/adapter.py index 8a7b129..58a79e9 100644 --- a/pyinfra/storage/adapters/adapter.py +++ b/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 diff --git a/pyinfra/storage/adapters/azure.py b/pyinfra/storage/adapters/azure.py index 04a8ca8..d0def0b 100644 --- a/pyinfra/storage/adapters/azure.py +++ b/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)) diff --git a/pyinfra/storage/adapters/s3.py b/pyinfra/storage/adapters/s3.py index f6bc110..93d4c0e 100644 --- a/pyinfra/storage/adapters/s3.py +++ b/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)) diff --git a/pyinfra/storage/storage.py b/pyinfra/storage/storage.py index 317da6e..4ed4822 100644 --- a/pyinfra/storage/storage.py +++ b/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 diff --git a/pyinfra/test/storage/adapter_mock.py b/pyinfra/test/storage/adapter_mock.py index 4406c5a..4550862 100644 --- a/pyinfra/test/storage/adapter_mock.py +++ b/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) diff --git a/pyinfra/test/storage/client_mock.py b/pyinfra/test/storage/client_mock.py index c0966a4..f6b6497 100644 --- a/pyinfra/test/storage/client_mock.py +++ b/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]) diff --git a/pyinfra/test/unit_tests/storage_test.py b/pyinfra/test/unit_tests/storage_test.py index b892c45..aacbd59 100644 --- a/pyinfra/test/unit_tests/storage_test.py +++ b/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": diff --git a/pytest.ini b/pytest.ini index 9b18a42..19deb86 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,5 +1,4 @@ [pytest] log_cli = 1 log_cli_level = DEBUG -addopts = --docker-compose=docker-compose.yaml diff --git a/scripts/mock_client.py b/scripts/mock_client.py index df9addc..d4be7cd 100644 --- a/scripts/mock_client.py +++ b/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}) diff --git a/src/serve.py b/src/serve.py index 5b0df96..a9d9e3b 100644 --- a/src/serve.py +++ b/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}'.")