Merge pull request #10 from kerthcet/feat/add-oss-s3

Support download_dir

Author: InftyAI-Agent

README.md

@@ -22,12 +22,24 @@ client = StoreFactory.new_client(
     provider=<provider>, endpoint=<endpoint>, bucket=<bucket>
 )
 
+# Create a directory
+client.create_dir(dir_name)
+
+# Delete a directory with all its files
+client.delete_dir(dir_name)
+
 # Upload
 client.upload(src, dest)
 
+# Upload a directory with all its files
+client.upload_dir(src_dir, dest_dir)
+
 # Download
 client.download(src, dest)
 
+# Download a directory with all its files
+client.download_dir(src_dir, dest_dir)
+
 # Exists
 client.exists(filename)
 

omnistore/objstore/aliyun_oss.py

@@ -1,3 +1,6 @@
+from operator import truediv
+from pathlib import Path
+
 import oss2
 from oss2.credentials import EnvironmentVariableCredentialsProvider
 
@@ -6,44 +9,53 @@
 
 class OSS(ObjStore):
     def __init__(self, endpoint: str, bucket: str):
-        """
-        Construct a new client to communicate with the provider.
-        """
-
         # Make sure environments OSS_ACCESS_KEY_ID and OSS_ACCESS_KEY_SECRET are exist.
         auth = oss2.ProviderAuth(EnvironmentVariableCredentialsProvider())
         self._bucket = oss2.Bucket(auth, endpoint, bucket)
 
-    def upload(self, src: str, dest: str):
-        """
-        Upload will upload the obj to the provider.
-        """
+    def create_dir(self, dirname: str):
+        self._bucket.put_object(dirname, "")
+
+    def delete_dir(self, dirname: str):
+        for obj in oss2.ObjectIterator(self._bucket, prefix=dirname):
+            self._bucket.delete_object(obj.key)
 
+    def upload(self, src: str, dest: str):
         oss2.resumable_upload(self._bucket, dest, src)
 
-    def download(self, src: str, dest: str):
-        """
-        Download will download the required obj from the provider.
-        """
+    def upload_dir(self, src_dir: str, dest_dir: str):
+        for file in Path(src_dir).rglob("*"):
+            if file.is_file():
+                self.upload(file, dest_dir + "/" + file.name)
+                continue
 
+            if file.is_dir():
+                # TODO: Support uploading subdirectory.
+                print("Don't support uploading subdirectory yet")
+
+    def download(self, src: str, dest: str):
         oss2.resumable_download(self._bucket, src, dest)
 
-    def delete(self, filename: str):
-        """
-        Delete will delete the obj from the provider.
-        """
+    def download_dir(self, src_dir: str, dest_dir: str):
+        if not src_dir.endswith("/"):
+            src_dir += "/"
+
+        if not dest_dir.endswith("/"):
+            dest_dir += "/"
+
+        path = Path(dest_dir)
+        if not path.exists():
+            path.mkdir(parents=True, exist_ok=True)
+
+        for obj in oss2.ObjectIterator(self._bucket, prefix=src_dir, delimiter="/"):
+            if obj.is_prefix():  # If this is a folder
+                # TODO: This is enough for download models, but not enough for general download usage.
+                print(f"Don't support downloading subdirectory: {obj.key} yet")
+            else:  # If this is a file
+                self.download(obj.key, dest_dir + obj.key.split("/")[-1])
 
+    def delete(self, filename: str):
         return self._bucket.delete_object(filename)
 
     def exists(self, filename: str):
-        """
-        Exists checks whether the obj exists in the provider.
-        """
-
         return self._bucket.object_exists(filename)
-
-    def accessible(self) -> bool:
-        """
-        Accessible checks whether the obj is visitable.
-        """
-        raise NotImplementedError("OSS not implemented")

omnistore/store.py

@@ -2,38 +2,58 @@
 
 
 class Store(ABC):
+    @abstractmethod
+    def create_dir(self, dirname: str):
+        """
+        create_dir will create a dir in the provider.
+        """
+        pass
+
+    @abstractmethod
+    def delete_dir(self, dirname: str):
+        """
+        delete_dir will delete a dir from the provider, together with all the files under the directory.
+        """
+        pass
+
     @abstractmethod
     def upload(self, src: str, dest: str):
         """
-        Upload will upload the obj to the provider.
+        upload will upload the file from src to dest, both parameters are file names.
+        """
+        pass
+
+    @abstractmethod
+    def upload_dir(self, src_dir: str, dest_dir: str):
+        """
+        upload_dir will upload the folder from src_dir to dest_dir, both parameters are folder names.
         """
         pass
 
     @abstractmethod
     def download(self, src: str, dest: str):
         """
-        download will download the required obj from the provider.
-        Non-null callback will be invocated when download finished.
+        download will download the file from src to dest, both parameters are file names.
         """
         pass
 
     @abstractmethod
-    def delete(self, filename: str):
+    def download_dir(self, src_dir: str, dest_dir: str):
         """
-        Delete will delete the obj from the provider.
+        download_dir will download the folder from src_dir to dest_dir, both parameters are folder names.
         """
         pass
 
     @abstractmethod
-    def exists(self, filename: str):
+    def delete(self, filename: str):
         """
-        Exists checks whether the obj exists in the provider.
+        delete will delete the file from the provider.
         """
         pass
 
     @abstractmethod
-    def accessible(self) -> bool:
+    def exists(self, filename: str):
         """
-        Accessible checks whether the obj is visitable.
+        exists checks whether the file exists in the provider.
         """
         pass

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnistore"
-version = "0.0.2"
+version = "0.0.3"
 description = "An unified python client to communicate with various kinds of object-store providers."
 authors = ["InftyAI Community"]
 license = "MIT"

tests/integration_tests/objstore/test_aliyun_oss.py

@@ -15,15 +15,12 @@ def setup_and_teardown(self):
         except Exception as e:
             print(f"An error occurred: {e}")
 
-        with open("./test-tmp/foo.txt", "w") as file:
-            file.write("test")
-
         yield
 
         print("Tearing down the test environment.")
         shutil.rmtree("./test-tmp")
 
-    def test_oss_operations(self):
+    def test_upload_and_download_files(self):
         endpoint = os.getenv("ENDPOINT")
         bucket = os.getenv("BUCKET")
 
@@ -32,6 +29,9 @@ def test_oss_operations(self):
         )
         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")
 
@@ -40,3 +40,25 @@ def test_oss_operations(self):
 
         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_OSS, endpoint=endpoint, bucket=bucket
+        )
+        assert False == client.exists("/test/foo.txt")
+
+        os.makedirs("./test-tmp/test", exist_ok=True)
+        with open("./test-tmp/test/foo.txt", "w") as file:
+            file.write("test")
+
+        client.upload_dir("./test-tmp/test", "test")
+        assert True == client.exists("test/foo.txt")
+
+        client.download_dir("test", "./test-tmp/test")
+        assert True == os.path.exists("./test-tmp/test/foo.txt")
+
+        client.delete_dir("test")
+        assert False == client.exists("test/foo.txt")