From a770acc6c4d1e50f142da913d4dc5ec23ed8d558 Mon Sep 17 00:00:00 2001 From: qinguoyi <1532979219@qq.com> Date: Mon, 21 Oct 2024 23:15:01 +0800 Subject: [PATCH] feat:add minio --- README.md | 17 ++++ omnistore/objstore/constant.py | 2 + omnistore/objstore/minio.py | 82 +++++++++++++++++++ omnistore/objstore/objstore_factory.py | 5 +- .../objstore/test_aliyun_oss.py | 2 + .../integration_tests/objstore/test_minio.py | 67 +++++++++++++++ 6 files changed, 173 insertions(+), 2 deletions(-) create mode 100644 omnistore/objstore/constant.py create mode 100644 omnistore/objstore/minio.py create mode 100644 tests/integration_tests/objstore/test_minio.py diff --git a/README.md b/README.md index 754e2b2..b2951d5 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,23 @@ OSS_ACCESS_KEY_ID= OSS_ACCESS_KEY_SECRET= ``` +### [MINIO](https://min.io/) + +Usage: + +```python +client = StoreFactory.new_client( + provider="MINIO", endpoint=, bucket= +) +``` + +Required environment variables: + +```yaml +MINIO_ACCESS_KEY= +MINIO_SECRET_KEY= +``` + ## Development Once you want to run the integration tests, you should have a `.env` file locally, similar to the `.env.example`. diff --git a/omnistore/objstore/constant.py b/omnistore/objstore/constant.py new file mode 100644 index 0000000..1e1d371 --- /dev/null +++ b/omnistore/objstore/constant.py @@ -0,0 +1,2 @@ +OBJECT_STORE_OSS = "OSS" +OBJECT_STORE_MINIO = "MINIO" diff --git a/omnistore/objstore/minio.py b/omnistore/objstore/minio.py new file mode 100644 index 0000000..31d6d2a --- /dev/null +++ b/omnistore/objstore/minio.py @@ -0,0 +1,82 @@ +import io +import os +from abc import ABC +from pathlib import Path + +import minio +from minio import credentials, S3Error + +from omnistore.objstore.objstore import ObjStore + + +class MinIO(ObjStore): + def __init__(self, endpoint: str, bucket: str): + """ + Construct a new client to communicate with the provider. + """ + auth = credentials.EnvMinioProvider() + self.client = minio.Minio(endpoint, credentials=auth,secure=False) + self.bucket_name = bucket + + # Make sure the bucket exists + if not self.client.bucket_exists(bucket): + self.client.make_bucket(bucket) + + def create_dir(self, dirname: str): + if not dirname.endswith("/"): + dirname += "/" + empty_stream = io.BytesIO(b"") + self.client.put_object(self.bucket_name, dirname, empty_stream, 0) + + def delete_dir(self, dirname: str): + if not dirname.endswith("/"): + dirname += "/" + objects = self.client.list_objects( + self.bucket_name, prefix=dirname, recursive=True + ) + for obj in objects: + self.client.remove_object(self.bucket_name, obj.object_name) + + def upload(self, src: str, dest: str): + self.client.fput_object(self.bucket_name, dest, src) + + def upload_dir(self, src_dir: str, dest_dir: str): + for file in Path(src_dir).rglob("*"): + if file.is_file(): + dest_path = f"{dest_dir}/{file.relative_to(src_dir)}" + self.upload(str(file), dest_path) + elif file.is_dir(): + self.create_dir(f"{dest_dir}/{file.relative_to(src_dir)}/") + + def download(self, src: str, dest: str): + self.client.fget_object(self.bucket_name, src, dest) + + def download_dir(self, src_dir: str, dest_dir: str): + if not src_dir.endswith("/"): + src_dir += "/" + path = Path(dest_dir) + if not path.exists(): + path.mkdir(parents=True) + objects = self.client.list_objects( + self.bucket_name, prefix=src_dir, recursive=True + ) + for obj in objects: + file_path = Path(dest_dir, Path(obj.object_name).relative_to(src_dir)) + if not file_path.parent.exists(): + file_path.parent.mkdir(parents=True, exist_ok=True) + if obj.object_name.endswith("/"): + continue + self.download(obj.object_name, str(file_path)) + + def delete(self, filename: str): + self.client.remove_object(self.bucket_name, filename) + + def exists(self, filename: str): + try: + self.client.stat_object(self.bucket_name, filename) + return True + except S3Error as e: + if e.code == "NoSuchKey": + return False + else: + raise e diff --git a/omnistore/objstore/objstore_factory.py b/omnistore/objstore/objstore_factory.py index ed2fa86..e2e60ba 100644 --- a/omnistore/objstore/objstore_factory.py +++ b/omnistore/objstore/objstore_factory.py @@ -1,12 +1,13 @@ from omnistore.objstore.aliyun_oss import OSS +from omnistore.objstore.constant import OBJECT_STORE_OSS, OBJECT_STORE_MINIO +from omnistore.objstore.minio import MinIO from omnistore.store import Store -OBJECT_STORE_OSS = "OSS" - class StoreFactory: ObjStores = { OBJECT_STORE_OSS: OSS, + OBJECT_STORE_MINIO: MinIO, } @classmethod diff --git a/tests/integration_tests/objstore/test_aliyun_oss.py b/tests/integration_tests/objstore/test_aliyun_oss.py index 587102a..734680a 100644 --- a/tests/integration_tests/objstore/test_aliyun_oss.py +++ b/tests/integration_tests/objstore/test_aliyun_oss.py @@ -2,9 +2,11 @@ import shutil import pytest +from dotenv import load_dotenv from omnistore.objstore.objstore_factory import OBJECT_STORE_OSS, StoreFactory +load_dotenv() class TestOSS: @pytest.fixture(scope="module", autouse=True) diff --git a/tests/integration_tests/objstore/test_minio.py b/tests/integration_tests/objstore/test_minio.py new file mode 100644 index 0000000..fe31ef2 --- /dev/null +++ b/tests/integration_tests/objstore/test_minio.py @@ -0,0 +1,67 @@ +import os +import shutil + +import pytest +from dotenv import load_dotenv + +from omnistore.objstore import StoreFactory +from omnistore.objstore.constant import OBJECT_STORE_MINIO + +load_dotenv() + +class TestMinio: + @pytest.fixture(scope="module", autouse=True) + def setup_and_teardown(self): + print("Setting up the test environment.") + try: + os.makedirs("./test-tmp", exist_ok=True) + except Exception as e: + print(f"An error occurred: {e}") + + yield + + print("Tearing down the test environment.") + shutil.rmtree("./test-tmp") + + def test_upload_and_download_files(self): + endpoint = os.getenv("ENDPOINT") + bucket = os.getenv("BUCKET") + + client = StoreFactory.new_client( + provider=OBJECT_STORE_MINIO, endpoint=endpoint, bucket=bucket + ) + assert False == client.exists("foo.txt") + + with open("./test-tmp/foo.txt", "w") as file: + file.write("test") + + client.upload("./test-tmp/foo.txt", "foo.txt") + assert True == client.exists("foo.txt") + + client.download("foo.txt", "./test-tmp/bar.txt") + assert True == os.path.exists("./test-tmp/bar.txt") + + client.delete("foo.txt") + assert False == client.exists("foo.txt") + + def test_upload_and_download_dir(self): + endpoint = os.getenv("ENDPOINT") + bucket = os.getenv("BUCKET") + + client = StoreFactory.new_client( + provider=OBJECT_STORE_MINIO, endpoint=endpoint, bucket=bucket + ) + assert False == client.exists("/test/foo.txt") + + os.makedirs("./test-tmp/test/111", exist_ok=True) + with open("./test-tmp/test/111/foo.txt", "w") as file: + file.write("test") + + client.upload_dir("./test-tmp/test", "test") + assert True == client.exists("test/111/foo.txt") + + client.download_dir("test", "./test-tmp/test1") + assert True == os.path.exists("./test-tmp/test1/111/foo.txt") + + client.delete_dir("test") + assert False == client.exists("test/foo.txt") \ No newline at end of file