49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
"""Definition of connect function."""
|
|
|
|
import os
|
|
from typing import Literal
|
|
|
|
from minio import Minio
|
|
|
|
|
|
def connect(
|
|
bucket: Literal['images', 'models'] = 'images',
|
|
) -> Minio:
|
|
"""Connect to MinIO server."""
|
|
assert isinstance(bucket, str)
|
|
assert bucket in ['images', 'models']
|
|
# ensure necessary env vars available
|
|
necesasary_var_list = [
|
|
'MINIO_ENDPOINT',
|
|
'MINIO_ACCESS_KEY',
|
|
'MINIO_SECRET_KEY',
|
|
'MINIO_BUCKET_NAME',
|
|
'MINIO_BUCKET_NAME_MODELS',
|
|
]
|
|
for env_var in necesasary_var_list:
|
|
# ensure env var set
|
|
assert (
|
|
env_var in os.environ
|
|
), f"environment variable not set: {
|
|
env_var
|
|
}"
|
|
# prepare arguments
|
|
minio_endpoint = os.getenv('MINIO_ENDPOINT')
|
|
minio_access_key = os.getenv('MINIO_ACCESS_KEY')
|
|
minio_secret_key = os.getenv('MINIO_SECRET_KEY')
|
|
if bucket == 'images':
|
|
minio_bucket_name = os.getenv('MINIO_BUCKET_NAME')
|
|
else:
|
|
minio_bucket_name = os.getenv('MINIO_BUCKET_NAME_MODELS')
|
|
# connect client
|
|
client = Minio(
|
|
endpoint=minio_endpoint,
|
|
access_key=minio_access_key,
|
|
secret_key=minio_secret_key,
|
|
secure=False,
|
|
)
|
|
# ensure bucket exists
|
|
if not client.bucket_exists(bucket_name=minio_bucket_name):
|
|
client.make_bucket(bucket_name=minio_bucket_name)
|
|
return client
|