Repository: AgnostiqHQ/covalent (python)

Issue:
File Transfer not working for Folders
### Environment

- **Covalent version**: 0.199.0rc0
- **Python version**: Python 3.8.13
- **Operating system**: macOS Monterey (v 12.5.1)


### What is happening?

I encountered the error **(404) when calling the HeadObject operation: Not Found** when attempting to transfer a folder from an S3 bucket to my local directory. Similar syntax does work with individual files, so this is an issue with transferring folders only.

### How can we reproduce the issue?

I defined the file transfer as below:

```
strategy = ct.fs_strategies.S3(credentials="~/.aws/credentials", profile="default", region_name="us-east-1")
from_remote_file = ct.fs.Folder("s3://my_data/folder", is_remote=True)
to_local_file = ct.fs.Folder("/My_local_dir/")
ft = ct.fs.FileTransfer(from_remote_file, to_local_file, strategy=strategy, order=ct.fs.Order.BEFORE)
```

And called it in the electron like this:

```
@ct.electron(files=[ft])
def preprocess_data(files=[]):
    y = "test"
    return y
```

with the lattice: 

```
@ct.lattice
def dummy_workflow(input):
    y = preprocess_data()
    return y
```

### What should happen?

The folder should be transferred from the S3 bucket to the specified local directory when the workflow is run.

### Any suggestions?

_No response_

Pull request:
Add folder upload/download to/from s3 buckets
Fixes #1316 
<!--
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Increment the version number in the /VERSION file. If this is a bugfix, increment the patch version (the rightmost number), for example 0.18.2 becomes 0.18.3. If it's a new feature, increment the minor version (the middle number), for example 0.18.2 becomes 0.19.0.
⚠️ Add a note to /CHANGELOG.md with your version number and the date, summarizing the changes.
⚠️ If your pull request fixes an open issue, please link to the issue.
⚠️ Rebase the latest changes from the develop branch. Assuming origin points to https://github.com/AgnostiqHQ/covalent, you should run git rebase origin/develop.
-->

- [ ] I have added the tests to cover my changes.
- [ ] I have updated the documentation and CHANGELOG accordingly.
- [ ] I have read the CONTRIBUTING document.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a17875bb..e6ba6955 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Added
+
+- Support for transferring the contents of folders to and from S3 buckets using the file transfer module.
+
 ## [0.207.0-rc.0] - 2022-10-26
 
 ### Authors
diff --git a/covalent/_file_transfer/strategies/s3_strategy.py b/covalent/_file_transfer/strategies/s3_strategy.py
index 37ed92bc..2ef6c5a3 100644
--- a/covalent/_file_transfer/strategies/s3_strategy.py
+++ b/covalent/_file_transfer/strategies/s3_strategy.py
@@ -69,42 +69,104 @@ class S3(FileTransferStrategy):
 
     # return callable to download here implies 'from' is a remote source
     def download(self, from_file: File, to_file: File = File()) -> File:
-        from_filepath = from_file.filepath.strip("/")
+        """Download files or the contents of folders from S3 bucket."""
+        app_log.debug(f"Is dir: {from_file._is_dir}")
+
+        from_filepath = from_file.filepath
         to_filepath = to_file.filepath
+
+        if from_filepath.startswith("/"):
+            from_filepath = from_filepath[1:]
+
         bucket_name = furl(from_file.uri).origin[5:]
         app_log.debug(
             f"S3 download bucket: {bucket_name}, from_filepath: {from_filepath}, to_filepath {to_filepath}."
         )
 
-        def callable():
-            import boto3
+        if from_file._is_dir:
 
-            s3 = boto3.client(
-                "s3",
-                region_name=self.region_name,
-            )
-            s3.download_file(bucket_name, from_filepath, to_filepath)
+            def callable():
+                """Download files from a folder in s3 bucket."""
+                from pathlib import Path
+
+                import boto3
+
+                s3 = boto3.client(
+                    "s3",
+                    region_name=self.region_name,
+                )
+                for obj_metadata in s3.list_objects(Bucket=bucket_name, Prefix=from_filepath)[
+                    "Contents"
+                ]:
+                    obj_key = obj_metadata["Key"]
+                    obj_destination_filepath = Path(to_filepath) / obj_key
+                    if obj_key.endswith("/"):
+                        obj_destination_filepath.mkdir(parents=True, exist_ok=True)
+                    else:
+                        s3.download_file(bucket_name, obj_key, str(obj_destination_filepath))
+
+        else:
+
+            def callable():
+                """Download file from s3 bucket."""
+                import boto3
+
+                s3 = boto3.client(
+                    "s3",
+                    region_name=self.region_name,
+                )
+                s3.download_file(bucket_name, from_filepath, to_filepath)
 
         return callable
 
     # return callable to download here implies 'to' is a remote source
     def upload(self, from_file: File, to_file: File = File()) -> File:
-
+        """Upload files or folders to S3 bucket."""
         from_filepath = from_file.filepath
-        to_filepath = to_file.filepath.strip("/")
+        to_filepath = to_file.filepath
+
+        if to_filepath.startswith("/"):
+            to_filepath = to_filepath.strip("/")
+
         bucket_name = furl(to_file.uri).origin[5:]
         app_log.debug(
-            f"S3 download bucket: {bucket_name}, from_filepath: {from_filepath}, to_filepath {to_filepath}."
+            f"S3 upload bucket: {bucket_name}, from_filepath: {from_filepath}, to_filepath {to_filepath}."
         )
 
-        def callable():
-            import boto3
+        if from_file._is_dir:
 
-            s3 = boto3.client(
-                "s3",
-                region_name=self.region_name,
-            )
-            s3.upload_file(from_filepath, bucket_name, to_filepath)
+            def callable():
+                """List and upload files from a directory to the remote S3 bucket."""
+                import os
+                from pathlib import Path
+
+                import boto3
+
+                s3 = boto3.client(
+                    "s3",
+                    region_name=self.region_name,
+                )
+
+                for dir_, _, files in os.walk(from_filepath):
+                    for file_name in files:
+                        rel_dir = os.path.relpath(dir_, from_filepath)
+                        rel_file = os.path.join(rel_dir, file_name)
+                        obj_from_filepath = str(Path(from_filepath) / rel_file)
+                        obj_to_filepath = str(Path(to_filepath) / rel_file)
+                        app_log.debug(f"Uploading from {obj_from_filepath} to {obj_to_filepath}.")
+                        s3.upload_file(obj_from_filepath, bucket_name, obj_to_filepath)
+
+        else:
+
+            def callable():
+                """Upload file to remote S3 bucket."""
+                import boto3
+
+                s3 = boto3.client(
+                    "s3",
+                    region_name=self.region_name,
+                )
+                s3.upload_file(from_filepath, bucket_name, to_filepath)
 
         return callable
 


Test patch:
diff --git a/tests/covalent_tests/file_transfer/strategies/s3_strategy_test.py b/tests/covalent_tests/file_transfer/strategies/s3_strategy_test.py
index 02644ea2..e7afdce4 100644
--- a/tests/covalent_tests/file_transfer/strategies/s3_strategy_test.py
+++ b/tests/covalent_tests/file_transfer/strategies/s3_strategy_test.py
@@ -1,3 +1,23 @@
+# Copyright 2021 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
 import os
 import sys
 from unittest.mock import MagicMock
@@ -5,7 +25,7 @@ from unittest.mock import MagicMock
 import pytest
 from furl import furl
 
-from covalent._file_transfer import File
+from covalent._file_transfer import File, Folder
 from covalent._file_transfer.strategies.s3_strategy import S3
 
 
@@ -84,3 +104,61 @@ class TestS3Strategy:
 
         with pytest.raises(NotImplementedError):
             S3().cp(File(self.MOCK_REMOTE_FILEPATH), File(self.MOCK_LOCAL_FILEPATH))()
+
+
+def test_folder_download(mocker):
+    """Test the s3 file download when remote and local folders are provided."""
+    boto3_mock = MagicMock()
+    sys.modules["boto3"] = boto3_mock
+
+    boto3_client_mock = mocker.patch("boto3.client")
+
+    from_folder = Folder("s3://mock-bucket/")
+    to_folder = Folder("/User/tmp/")
+
+    bucket_name = furl(from_folder.uri).origin[5:]
+
+    boto3_client_mock().list_objects.return_value = {
+        "Contents": [
+            {"Key": "test.csv"},
+        ]
+    }
+
+    callable_func = S3().download(from_folder, to_folder)
+    callable_func()
+    boto3_client_mock().download_file.assert_called_once_with(
+        bucket_name, "test.csv", "/User/tmp/test.csv"
+    )
+
+
+def test_folder_upload(mocker):
+    """Test the s3 file upload method when remote and local folders are provided."""
+    boto3_mock = MagicMock()
+    sys.modules["boto3"] = boto3_mock
+
+    boto3_client_mock = mocker.patch("boto3.client")
+
+    os_mock = MagicMock()
+    sys.modules["os"] = os_mock
+
+    mocker.patch("os.walk", return_value=[["", "", ["test.csv"]]])
+
+    to_folder = Folder("s3://mock-bucket/")
+    from_folder = Folder("/User/tmp/")
+
+    bucket_name = furl(to_folder.uri).origin[5:]
+
+    boto3_client_mock().list_objects.return_value = {
+        "Contents": [
+            {"Key": "test.csv"},
+        ]
+    }
+
+    os_mock.path.relpath.return_value = "mock_path"
+    os_mock.path.join.return_value = "mock_join"
+
+    callable_func = S3().upload(from_folder, to_folder)
+    callable_func()
+    boto3_client_mock().upload_file.assert_called_once_with(
+        "/User/tmp/mock_join", bucket_name, "mock_join"
+    )