Skip to content

Commit

Permalink
feat:add minio
Browse files Browse the repository at this point in the history
  • Loading branch information
qinguoyi committed Oct 23, 2024
1 parent 81c4978 commit a770acc
Show file tree
Hide file tree
Showing 6 changed files with 173 additions and 2 deletions.
17 changes: 17 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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=<endpoint>, bucket=<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`.
Expand Down
2 changes: 2 additions & 0 deletions omnistore/objstore/constant.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
OBJECT_STORE_OSS = "OSS"
OBJECT_STORE_MINIO = "MINIO"
82 changes: 82 additions & 0 deletions omnistore/objstore/minio.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 3 additions & 2 deletions omnistore/objstore/objstore_factory.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 2 additions & 0 deletions tests/integration_tests/objstore/test_aliyun_oss.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
67 changes: 67 additions & 0 deletions tests/integration_tests/objstore/test_minio.py
Original file line number Diff line number Diff line change
@@ -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")

0 comments on commit a770acc

Please sign in to comment.