model/src/data_store/connect.py

@@ -1,30 +0,0 @@
-"""Definition of connect function."""
-from __future__ import annotations
-
-import os
-
-from minio import Minio
-
-
-def connect(
-) -> Minio:
-    """Connect to MinIO server"""
-    minio_endpoint = os.getenv('MINIO_ENDPOINT', default=None)
-    assert isinstance(minio_endpoint, str)
-    minio_access_key = os.getenv('MINIO_ACCESS_KEY', default=None)
-    assert isinstance(minio_access_key, str)
-    minio_secret_key = os.getenv('MINIO_SECRET_KEY', default=None)
-    assert isinstance(minio_secret_key, str)
-    minio_bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
-    assert isinstance(minio_bucket_name, str)
-    # 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

shared/data_store/connect.py

@@ -1,37 +1,40 @@
 """Definition of connect function."""
-from __future__ import annotations
 
 import os
+from typing import Literal
 
 from minio import Minio
 
 
 def connect(
+    bucket: Literal['images', 'models'] = 'images',
 ) -> Minio:
-    """Connect to MinIO server"""
+    """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}"
-        )
+        ), f"environment variable not set: {
+            env_var
+        }"
     # prepare arguments
-    minio_endpoint = os.getenv('MINIO_ENDPOINT', default=None)
-    assert isinstance(minio_endpoint, str)
-    minio_access_key = os.getenv('MINIO_ACCESS_KEY', default=None)
-    assert isinstance(minio_access_key, str)
-    minio_secret_key = os.getenv('MINIO_SECRET_KEY', default=None)
-    assert isinstance(minio_secret_key, str)
-    minio_bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
-    assert isinstance(minio_bucket_name, str)
+    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,