refactoring; implemented S3 handle

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:
+s3:
-  endpoint: $STORAGE_ENDPOINT|"127.0.0.1:9000"  # MinIO endpoint
+  endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000"  # MinIO endpoint
-  user: $STORAGE_KEY|root  # MinIO user name
+  access_key: $STORAGE_KEY|root  # MinIO user name
-  password: $STORAGE_SECRET|password  # MinIO user password
+  secret_key: $STORAGE_SECRET|password  # MinIO user password
   bucket: $STORAGE_BUCKET_NAME|redaction  # MinIO bucket
 
 azure_blob_storage:

pyinfra/exceptions.py

@@ -13,5 +13,10 @@ class ProcessingFailure(Exception):
 class UnknownStorageBackend(ValueError):
     pass
 
+
 class InvalidEndpoint(ValueError):
     pass
+
+
+class UnknownClient(ValueError):
+    pass

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

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:
+    def bucket_exists(self, bucket_name):
-        container_client = self.__client.get_container_client(container_name)
+        container_client = self.__client.get_container_client(bucket_name)
-        return container_client if container_client.exists() else self.__client.create_container(container_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}'...")

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)

pyinfra/storage/clients/s3.py

@@ -7,7 +7,7 @@ from pyinfra.exceptions import InvalidEndpoint
 
 
 def parse_endpoint(endpoint):
-    endpoint_pattern = r"(?P<protocol>https?)://(?P<address>.*:(?:\d{1,3}\.){3}\d{1,3})"
+    endpoint_pattern = r"(?P<protocol>https?)://(?P<address>(?:(?:(?:\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)

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)
 

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)
 

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

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