pyinfra/pyinfra/storage/adapters/s3.py

import io
from itertools import repeat
import logging
from operator import attrgetter

from minio import Minio

from pyinfra.storage.adapters.adapter import StorageAdapter

logger = logging.getLogger(__name__)


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.has_bucket(bucket_name):
            self.__client.make_bucket(bucket_name)

    def has_bucket(self, bucket_name):
        return self.__client.bucket_exists(bucket_name)

    def put_object(self, bucket_name, object_name, data):
        logger.debug(f"Uploading '{object_name}'...")
        data = io.BytesIO(data)
        self.__client.put_object(bucket_name, object_name, data, length=data.getbuffer().nbytes)

    def get_object(self, bucket_name, object_name):
        logger.debug(f"Downloading '{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_objects(self, bucket_name):
        for obj in self.__client.list_objects(bucket_name, recursive=True):
            logger.debug(f"Downloading '{obj.object_name}'...")
            yield self.get_object(bucket_name, obj.object_name)

    def clear_bucket(self, bucket_name):
        logger.debug(f"Clearing S3 bucket '{bucket_name}'...")
        objects = self.__client.list_objects(bucket_name, recursive=True)
        for obj in objects:
            self.__client.remove_object(bucket_name, obj.object_name)

    def get_all_object_names(self, bucket_name):
        objs = self.__client.list_objects(bucket_name, recursive=True)
        return zip(repeat(bucket_name), map(attrgetter("object_name"), objs))