80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
import io
|
|
import logging
|
|
from itertools import repeat
|
|
from operator import attrgetter
|
|
from urllib.parse import urlparse
|
|
|
|
from minio import Minio
|
|
from retry import retry
|
|
from validators import url
|
|
|
|
from pyinfra.config import Config, get_config
|
|
|
|
CONFIG = get_config()
|
|
logger = logging.getLogger(CONFIG.logging_level_root)
|
|
|
|
|
|
class S3StorageAdapter(object):
|
|
def __init__(self, client: Minio):
|
|
self._client = 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)
|
|
|
|
@retry(tries=3, delay=5, jitter=(1, 3))
|
|
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 Exception("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):
|
|
objs = self._client.list_objects(bucket_name, recursive=True)
|
|
return zip(repeat(bucket_name), map(attrgetter("object_name"), objs))
|
|
|
|
|
|
def _parse_endpoint(endpoint):
|
|
if url(endpoint):
|
|
parsed_url = urlparse(endpoint)
|
|
return {"secure": parsed_url.scheme == "https", "endpoint": parsed_url.netloc}
|
|
else:
|
|
raise Exception(f"The configured storage endpoint is not a valid url: {endpoint}")
|
|
|
|
|
|
def get_s3_storage(config: Config):
|
|
return S3StorageAdapter(Minio(
|
|
**_parse_endpoint(config.storage_endpoint),
|
|
access_key=config.storage_key,
|
|
secret_key=config.storage_secret,
|
|
# FIXME Is this still needed? Check and if yes, add it to config
|
|
# region=config.region,
|
|
))
|