From 16e34052da2cc8c61473d21c1929023445e619d6 Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Thu, 24 Feb 2022 14:11:53 +0100 Subject: [PATCH] refactoring; implemented S3 handle --- config.yaml | 8 ++--- pyinfra/exceptions.py | 5 +++ pyinfra/storage/adapters/adapter.py | 8 +++++ pyinfra/storage/adapters/azure.py | 15 +++++++-- pyinfra/storage/adapters/s3.py | 43 +++++++++++++++++++++++++ pyinfra/storage/clients/s3.py | 4 +-- pyinfra/storage/storage.py | 6 ++++ pyinfra/test/storage/adapter_mock.py | 6 ++++ pyinfra/test/storage/client_mock.py | 11 ++++--- pyinfra/test/unit_tests/storage_test.py | 14 +++++++- 10 files changed, 106 insertions(+), 14 deletions(-) diff --git a/config.yaml b/config.yaml index 5cbe777..052d4c8 100755 --- a/config.yaml +++ b/config.yaml @@ -16,10 +16,10 @@ rabbitmq: enabled: $RETRY|False # Toggles retry behaviour max_attempts: $MAX_ATTEMPTS|3 # Number of times a message may fail before being published to dead letter queue -minio: - endpoint: $STORAGE_ENDPOINT|"127.0.0.1:9000" # MinIO endpoint - user: $STORAGE_KEY|root # MinIO user name - password: $STORAGE_SECRET|password # MinIO user password +s3: + endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000" # MinIO endpoint + access_key: $STORAGE_KEY|root # MinIO user name + secret_key: $STORAGE_SECRET|password # MinIO user password bucket: $STORAGE_BUCKET_NAME|redaction # MinIO bucket azure_blob_storage: diff --git a/pyinfra/exceptions.py b/pyinfra/exceptions.py index ab3a89b..8c66c79 100644 --- a/pyinfra/exceptions.py +++ b/pyinfra/exceptions.py @@ -13,5 +13,10 @@ class ProcessingFailure(Exception): class UnknownStorageBackend(ValueError): pass + class InvalidEndpoint(ValueError): pass + + +class UnknownClient(ValueError): + pass diff --git a/pyinfra/storage/adapters/adapter.py b/pyinfra/storage/adapters/adapter.py index 61b51e8..7d86631 100644 --- a/pyinfra/storage/adapters/adapter.py +++ b/pyinfra/storage/adapters/adapter.py @@ -5,6 +5,14 @@ class StorageAdapter(ABC): def __init__(self, client): self.__client = client + @abstractmethod + def make_bucket(self, bucket_name): + pass + + @abstractmethod + def bucket_exists(self, bucket_name): + pass + @abstractmethod def put(self, bucket_name, object_name, data): pass diff --git a/pyinfra/storage/adapters/azure.py b/pyinfra/storage/adapters/azure.py index 5566584..8092834 100644 --- a/pyinfra/storage/adapters/azure.py +++ b/pyinfra/storage/adapters/azure.py @@ -31,9 +31,18 @@ class AzureStorageAdapter(StorageAdapter): super().__init__(client=client) self.__client: BlobServiceClient = self._StorageAdapter__client - def __provide_container_client(self, container_name) -> ContainerClient: - container_client = self.__client.get_container_client(container_name) - return container_client if container_client.exists() else self.__client.create_container(container_name) + def bucket_exists(self, bucket_name): + container_client = self.__client.get_container_client(bucket_name) + return container_client.exists() + + def make_bucket(self, bucket_name): + container_client = self.__client.get_container_client(bucket_name) + container_client if container_client.exists() else self.__client.create_container(bucket_name) + + def __provide_container_client(self, bucket_name) -> ContainerClient: + self.make_bucket(bucket_name) + container_client = self.__client.get_container_client(bucket_name) + return container_client def put(self, bucket_name, object_name, data): logger.debug(f"Uploading '{object_name}'...") diff --git a/pyinfra/storage/adapters/s3.py b/pyinfra/storage/adapters/s3.py index e69de29..8faf336 100644 --- a/pyinfra/storage/adapters/s3.py +++ b/pyinfra/storage/adapters/s3.py @@ -0,0 +1,43 @@ +import io + +from minio import Minio + +from pyinfra.storage.adapters.adapter import StorageAdapter + + +class S3StorageAdapter(StorageAdapter): + + def __init__(self, client): + super().__init__(client=client) + self.__client: Minio = self._StorageAdapter__client + + def make_bucket(self, bucket_name): + if not self.bucket_exists(bucket_name): + self.__client.make_bucket(bucket_name) + + def bucket_exists(self, bucket_name): + return self.__client.bucket_exists(bucket_name) + + def put(self, bucket_name, object_name, data): + data = io.BytesIO(data) + self.__client.put_object(bucket_name, object_name, data, length=data.getbuffer().nbytes) + + def get(self, bucket_name, object_name): + response = None + + try: + response = self.__client.get_object(bucket_name, object_name) + return response.data + finally: + if response: + response.close() + response.release_conn() + + def get_all(self, bucket_name): + for obj in self.__client.list_objects(bucket_name): + yield self.get(bucket_name, obj.object_name) + + def purge(self, bucket_name): + objects = self.__client.list_objects(bucket_name) + for obj in objects: + self.__client.remove_object(bucket_name, obj.object_name) diff --git a/pyinfra/storage/clients/s3.py b/pyinfra/storage/clients/s3.py index 05f86ad..10abdd3 100644 --- a/pyinfra/storage/clients/s3.py +++ b/pyinfra/storage/clients/s3.py @@ -7,7 +7,7 @@ from pyinfra.exceptions import InvalidEndpoint def parse_endpoint(endpoint): - endpoint_pattern = r"(?Phttps?)://(?P
.*:(?:\d{1,3}\.){3}\d{1,3})" + endpoint_pattern = r"(?Phttps?)://(?P
(?:(?:(?:\d{1,3}\.){3}\d{1,3})|\w+):\d+)" match = re.match(endpoint_pattern, endpoint) @@ -19,4 +19,4 @@ def parse_endpoint(endpoint): def get_s3_client(endpoint=None) -> Minio: endpoint = CONFIG.s3.endpoint if endpoint is None else endpoint - return Minio(**parse_endpoint(endpoint), access_key=CONFIG.s3.secret_key, secret_key=CONFIG.s3.access_key) + return Minio(**parse_endpoint(endpoint), access_key=CONFIG.s3.access_key, secret_key=CONFIG.s3.secret_key) diff --git a/pyinfra/storage/storage.py b/pyinfra/storage/storage.py index cabe6c5..0177dc9 100644 --- a/pyinfra/storage/storage.py +++ b/pyinfra/storage/storage.py @@ -2,6 +2,12 @@ class Storage: def __init__(self, adapter): self.__adapter = adapter + def make_bucket(self, bucket_name): + self.__adapter.make_bucket(bucket_name) + + def bucket_exists(self, bucket_name): + return self.__adapter.bucket_exists(bucket_name) + def put(self, bucket_name, object_name, data): self.__adapter.put(bucket_name, object_name, data) diff --git a/pyinfra/test/storage/adapter_mock.py b/pyinfra/test/storage/adapter_mock.py index 467dac3..19e4720 100644 --- a/pyinfra/test/storage/adapter_mock.py +++ b/pyinfra/test/storage/adapter_mock.py @@ -8,6 +8,12 @@ class StorageAdapterMock(StorageAdapter): super().__init__(client=client) self.__client = self._StorageAdapter__client + def make_bucket(self, bucket_name): + self.__client.make_bucket(bucket_name) + + def bucket_exists(self, bucket_name): + return self.__client.bucket_exists(bucket_name) + def put(self, bucket_name, object_name, data): return self.__client.put(bucket_name, object_name, data) diff --git a/pyinfra/test/storage/client_mock.py b/pyinfra/test/storage/client_mock.py index cfc3eaa..2b2fb26 100644 --- a/pyinfra/test/storage/client_mock.py +++ b/pyinfra/test/storage/client_mock.py @@ -1,9 +1,12 @@ -from collections import defaultdict - - class StorageClientMock: def __init__(self): - self.__data = defaultdict(dict) + self.__data = {} + + def make_bucket(self, bucket_name): + self.__data[bucket_name] = {} + + def bucket_exists(self, bucket_name): + return bucket_name in self.__data def put(self, bucket_name, object_name, data): self.__data[bucket_name][object_name] = data diff --git a/pyinfra/test/unit_tests/storage_test.py b/pyinfra/test/unit_tests/storage_test.py index 38ec3b3..712ae7b 100644 --- a/pyinfra/test/unit_tests/storage_test.py +++ b/pyinfra/test/unit_tests/storage_test.py @@ -2,14 +2,17 @@ import logging import pytest +from pyinfra.exceptions import UnknownClient from pyinfra.storage.adapters.azure import AzureStorageAdapter +from pyinfra.storage.adapters.s3 import S3StorageAdapter from pyinfra.storage.clients.azure import get_azure_client +from pyinfra.storage.clients.s3 import get_s3_client from pyinfra.storage.storage import Storage from pyinfra.test.storage.adapter_mock import StorageAdapterMock from pyinfra.test.storage.client_mock import StorageClientMock -@pytest.mark.parametrize("client_name", ["mock", "azure"]) +@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"]) class TestStorage: def test_purging_bucket_yields_empty_bucket(self, storage, bucket_name): storage.purge(bucket_name) @@ -27,12 +30,20 @@ class TestStorage: data_received = storage.get_all(bucket_name) assert {b"content 1", b"content 2"} == {*data_received} + def test_make_bucket_produces_bucket(self, storage, bucket_name): + storage.make_bucket(bucket_name) + assert storage.bucket_exists(bucket_name) + def get_adapter(client_name): if client_name == "mock": return StorageAdapterMock(StorageClientMock()) if client_name == "azure": return AzureStorageAdapter(get_azure_client()) + if client_name == "s3": + return S3StorageAdapter(get_s3_client()) + else: + raise UnknownClient(client_name) @pytest.fixture @@ -44,6 +55,7 @@ def adapter(client_name): @pytest.fixture def storage(client_name, bucket_name): storage = Storage(get_adapter(client_name)) + storage.make_bucket(bucket_name) storage.purge(bucket_name) logging.info(client_name) return storage