60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
import io
|
|
import logging
|
|
from itertools import repeat
|
|
from operator import attrgetter
|
|
|
|
from minio import Minio
|
|
|
|
from pyinfra.exceptions import DataLoadingFailure
|
|
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)
|
|
# input("uploaded")
|
|
|
|
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
|
|
except Exception as err:
|
|
raise DataLoadingFailure("Failed getting object from s3 client") from err
|
|
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, prefix=None):
|
|
objs = self.__client.list_objects(bucket_name, recursive=True, prefix=prefix)
|
|
return map(attrgetter("object_name"), objs)
|