From 74d73072194e1d4a76e1ff54f2e0e5e8f247d81a Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Sat, 13 Jul 2024 22:36:45 +0200
Subject: [PATCH] moved to shared folder

---
 model/src/data_store/connect.py | 30 ------------------------------
 shared/data_store/connect.py    | 29 ++++++++++++++++-------------
 2 files changed, 16 insertions(+), 43 deletions(-)
 delete mode 100644 model/src/data_store/connect.py

diff --git a/model/src/data_store/connect.py b/model/src/data_store/connect.py
deleted file mode 100644
index 0bbf3af..0000000
--- a/model/src/data_store/connect.py
+++ /dev/null
@@ -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
diff --git a/shared/data_store/connect.py b/shared/data_store/connect.py
index 6530e78..7e102e8 100644
--- a/shared/data_store/connect.py
+++ b/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,