Repository: AgnostiqHQ/covalent
Pull request #1389 (2022-10-26 23:46:07)

S3 file transfer compatibility with profile, and region in remote execution backends
- [ ] I have added the tests to cover my changes.
- [ ] I have updated the documentation and CHANGELOG accordingly.
- [ ] I have read the CONTRIBUTING document.


Fixes two issues:
- If S3 Strategy is used in a remote executor profile & region were not being propagated properly, now they are, and if not specified it uses the AWS profile & region default env vars defined in the **execution backend**
- If S3 Strategy is used in a remote executor the validate credentials step in the S3 file transfer strategy would run locally, hence may erroneously assert that credentials are not valid as the credentials specified were intended to exist only in the execution backend

codecov[bot]: # [Codecov](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1389?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) Report
> Merging [#1389](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1389?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (e047dd5) into [develop](https://codecov.io/gh/AgnostiqHQ/covalent/commit/a6a7ca2edc433246f238d06c1bfd920953caca2d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (a6a7ca2) will **increase** coverage by `6.70%`.
> The diff coverage is `92.30%`.

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1389/graphs/tree.svg?width=650&height=150&src=pr&token=YGHCB3DE4P&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1389?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)

```diff
@@             Coverage Diff             @@
##           develop    #1389      +/-   ##
===========================================
+ Coverage    76.13%   82.84%   +6.70%     
===========================================
  Files          141       65      -76     
  Lines         5574     3900    -1674     
  Branches       145        0     -145     
===========================================
- Hits          4244     3231    -1013     
+ Misses        1208      669     -539     
+ Partials       122        0     -122     
```

| Flag | Coverage Δ | | *Carryforward flag |
|---|---|---|---|
| Dispatcher | `86.14% <ø> (ø)` | | Carriedforward from [a6a7ca2](https://codecov.io/gh/AgnostiqHQ/covalent/commit/a6a7ca2edc433246f238d06c1bfd920953caca2d?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) |
| Functional_Tests | `?` | | |
| SDK | `80.74% <92.30%> (-0.07%)` | :arrow_down: | |
| UI_Frontend | `?` | | |

*This pull request uses carry forward flags. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) to find out more.

</details>

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 14e58378..fed49512 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]
 
+### Fixed
+
+- Respecting specified AWS profile & region in remote executed S3 file transfers, defaulting to env vars of execution backend
+
 ### Added
 
 - Added `TaskRuntimeError` exception for executor plugin implementations to signal to Covalent that a task raised an
diff --git a/covalent/_file_transfer/strategies/s3_strategy.py b/covalent/_file_transfer/strategies/s3_strategy.py
index ed84b65c..9a19b2ca 100644
--- a/covalent/_file_transfer/strategies/s3_strategy.py
+++ b/covalent/_file_transfer/strategies/s3_strategy.py
@@ -36,13 +36,7 @@ class S3(FileTransferStrategy):
     Implements Base FileTransferStrategy class to upload/download files from S3 Bucket.
     """
 
-    def __init__(
-        self,
-        credentials: str = os.environ.get("AWS_SHARED_CREDENTIALS_FILE")
-        or os.path.join(os.environ["HOME"], ".aws/credentials"),
-        profile: str = os.environ.get("AWS_PROFILE") or None,
-        region_name: str = os.getenv("AWS_REGION") or None,
-    ):
+    def __init__(self, credentials: str = None, profile: str = None, region_name: str = None):
 
         self.credentials = credentials
         self.profile = profile
@@ -55,17 +49,8 @@ class S3(FileTransferStrategy):
                 "Using S3 strategy requires boto3 from AWS installed on your system."
             )
 
-        os.environ["AWS_SHARED_CREDENTIALS_FILE"] = self.credentials
-        if self.profile is not None:
-            os.environ["AWS_PROFILE"] = self.profile
-
-        # AWS Account Retrieval
-        sts = boto3.client("sts")
-        identity = sts.get_caller_identity()
-        account = identity.get("Account")
-
-        if account is None:
-            raise Exception("Incorrect AWS account credentials")
+        if self.credentials:
+            os.environ["AWS_SHARED_CREDENTIALS_FILE"] = self.credentials
 
     # return callable to download here implies 'from' is a remote source
     def download(self, from_file: File, to_file: File = File()) -> File:
@@ -83,6 +68,17 @@ class S3(FileTransferStrategy):
             f"S3 download bucket: {bucket_name}, from_filepath: {from_filepath}, to_filepath {to_filepath}."
         )
 
+        executor_profile = self.profile
+        executor_region = self.region_name
+
+        def get_boto_options(profile=None, region=None):
+            boto_options = {}
+            if profile:
+                boto_options["profile_name"] = profile
+            if region:
+                boto_options["region_name"] = region
+            return boto_options
+
         if from_file._is_dir:
 
             def callable():
@@ -91,10 +87,9 @@ class S3(FileTransferStrategy):
 
                 import boto3
 
-                s3 = boto3.client(
-                    "s3",
-                    region_name=self.region_name,
-                )
+                profile = executor_profile
+                region = executor_region
+                s3 = boto3.Session(**get_boto_options(profile, region)).client("s3")
 
                 for obj_metadata in s3.list_objects(Bucket=bucket_name, Prefix=from_filepath)[
                     "Contents"
@@ -118,10 +113,10 @@ class S3(FileTransferStrategy):
                 """Download file from s3 bucket."""
                 import boto3
 
-                s3 = boto3.client(
-                    "s3",
-                    region_name=self.region_name,
-                )
+                profile = executor_profile
+                region = executor_region
+                s3 = boto3.Session(**get_boto_options(profile, region)).client("s3")
+
                 s3.download_file(bucket_name, from_filepath, to_filepath)
 
         return callable
@@ -140,6 +135,17 @@ class S3(FileTransferStrategy):
             f"S3 upload bucket: {bucket_name}, from_filepath: {from_filepath}, to_filepath {to_filepath}."
         )
 
+        executor_profile = self.profile
+        executor_region = self.region_name
+
+        def get_boto_options(profile=None, region=None):
+            boto_options = {}
+            if profile:
+                boto_options["profile_name"] = profile
+            if region:
+                boto_options["region_name"] = region
+            return boto_options
+
         if from_file._is_dir:
 
             def callable():
@@ -149,10 +155,9 @@ class S3(FileTransferStrategy):
 
                 import boto3
 
-                s3 = boto3.client(
-                    "s3",
-                    region_name=self.region_name,
-                )
+                profile = executor_profile
+                region = executor_region
+                s3 = boto3.Session(**get_boto_options(profile, region)).client("s3")
 
                 for dir_, _, files in os.walk(from_filepath):
                     for file_name in files:
@@ -169,10 +174,10 @@ class S3(FileTransferStrategy):
                 """Upload file to remote S3 bucket."""
                 import boto3
 
-                s3 = boto3.client(
-                    "s3",
-                    region_name=self.region_name,
-                )
+                profile = executor_profile
+                region = executor_region
+                s3 = boto3.Session(**get_boto_options(profile, region)).client("s3")
+
                 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 7cc64321..b94b2df4 100644
--- a/tests/covalent_tests/file_transfer/strategies/s3_strategy_test.py
+++ b/tests/covalent_tests/file_transfer/strategies/s3_strategy_test.py
@@ -18,7 +18,6 @@
 #
 # Relief from the License may be granted by purchasing a commercial license.
 
-import os
 import sys
 from unittest.mock import MagicMock
 
@@ -34,28 +33,30 @@ class TestS3Strategy:
     MOCK_LOCAL_FILEPATH = "/Users/user/data.csv"
     MOCK_REMOTE_FILEPATH = "s3://covalent-tmp/data.csv"
     MOCK_AWS_CREDENTIALS = "~/.aws/credentials"
-    MOCK_AWS_PROFILE = "default"
-    MOCK_AWS_REGION = "us-east-1"
+    MOCK_AWS_PROFILE = "my_profile"
+    MOCK_AWS_REGION = "some-region-1"
+
+    @property
+    def MOCK_STRATEGY_CONFIG(self):
+        return {
+            "profile": self.MOCK_AWS_PROFILE,
+            "region_name": self.MOCK_AWS_REGION,
+            "credentials": self.MOCK_AWS_CREDENTIALS,
+        }
 
     def test_init(self, mocker):
-        # test S3 init fails due to not having boto3
-        with pytest.raises(ImportError):
-            strategy = S3()
 
-        # test with boto3 module defined and failure resolving aws credentials
         boto3_mock = MagicMock()
         sys.modules["boto3"] = boto3_mock
         boto3_client_mock = mocker.patch("boto3.client")
-        boto3_client_mock().get_caller_identity().get.return_value = None
-        with pytest.raises(Exception):
-            strategy = S3()
 
         # test with boto3 module defined and successfully resolving aws credentials - success initating S3 strategy
         account_mock = MagicMock()
         boto3_client_mock().get_caller_identity().get.return_value = account_mock
         strategy = S3(self.MOCK_AWS_CREDENTIALS, self.MOCK_AWS_PROFILE, self.MOCK_AWS_REGION)
-        assert os.environ["AWS_SHARED_CREDENTIALS_FILE"] == self.MOCK_AWS_CREDENTIALS
-        assert os.environ["AWS_PROFILE"] == self.MOCK_AWS_PROFILE
+        assert strategy.credentials == self.MOCK_AWS_CREDENTIALS
+        assert strategy.profile == self.MOCK_AWS_PROFILE
+        assert strategy.region_name == self.MOCK_AWS_REGION
 
     def test_download(self, mocker):
         # validate boto3.client('s3').download_file is called with appropriate arguments
@@ -63,7 +64,7 @@ class TestS3Strategy:
         boto3_mock = MagicMock()
         sys.modules["boto3"] = boto3_mock
 
-        boto3_client_mock = mocker.patch("boto3.client")
+        boto3_client_mock = boto3_mock.Session().client
         boto3_client_mock.download_file.return_value = None
 
         from_file = File(self.MOCK_REMOTE_FILEPATH)
@@ -71,7 +72,7 @@ class TestS3Strategy:
 
         bucket_name = furl(from_file.uri).origin[5:]
 
-        S3().download(from_file, to_file)()
+        S3(**self.MOCK_STRATEGY_CONFIG).download(from_file, to_file)()
 
         boto3_client_mock().download_file.assert_called_with(
             bucket_name, from_file.filepath.strip("/"), to_file.filepath
@@ -83,7 +84,7 @@ class TestS3Strategy:
         boto3_mock = MagicMock()
         sys.modules["boto3"] = boto3_mock
 
-        boto3_client_mock = mocker.patch("boto3.client")
+        boto3_client_mock = boto3_mock.Session().client
         boto3_client_mock.download_file.return_value = None
 
         to_file = File(self.MOCK_REMOTE_FILEPATH)
@@ -91,7 +92,7 @@ class TestS3Strategy:
 
         bucket_name = furl(to_file.uri).origin[5:]
 
-        S3().upload(from_file, to_file)()
+        S3(**self.MOCK_STRATEGY_CONFIG).upload(from_file, to_file)()
 
         boto3_client_mock().upload_file.assert_called_with(
             from_file.filepath, bucket_name, to_file.filepath.strip("/")
@@ -105,60 +106,57 @@ class TestS3Strategy:
         with pytest.raises(NotImplementedError):
             S3().cp(File(self.MOCK_REMOTE_FILEPATH), File(self.MOCK_LOCAL_FILEPATH))()
 
+    def test_folder_download(self, mocker):
+        """Test the s3 file download when remote and local folders are provided."""
+        boto3_mock = MagicMock()
+        sys.modules["boto3"] = boto3_mock
 
-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("/tmp/")
-
-    bucket_name = furl(from_folder.uri).origin[5:]
+        boto3_client_mock = boto3_mock.Session().client
 
-    boto3_client_mock().list_objects.return_value = {
-        "Contents": [
-            {"Key": "test.csv"},
-        ]
-    }
+        from_folder = Folder("s3://mock-bucket/")
+        to_folder = Folder("/tmp/")
 
-    callable_func = S3().download(from_folder, to_folder)
-    callable_func()
-    boto3_client_mock().download_file.assert_called_once_with(
-        bucket_name, "test.csv", "/tmp/test.csv"
-    )
+        bucket_name = furl(from_folder.uri).origin[5:]
 
+        boto3_client_mock().list_objects.return_value = {
+            "Contents": [
+                {"Key": "test.csv"},
+            ]
+        }
+        callable_func = S3(**self.MOCK_STRATEGY_CONFIG).download(from_folder, to_folder)
+        callable_func()
+        boto3_client_mock().download_file.assert_called_once_with(
+            bucket_name, "test.csv", "/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
+    def test_folder_upload(self, 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")
+        boto3_client_mock = boto3_mock.Session().client
 
-    os_mock = MagicMock()
-    sys.modules["os"] = os_mock
+        os_mock = MagicMock()
+        sys.modules["os"] = os_mock
 
-    mocker.patch("os.walk", return_value=[["", "", ["test.csv"]]])
+        mocker.patch("os.walk", return_value=[["", "", ["test.csv"]]])
 
-    to_folder = Folder("s3://mock-bucket/")
-    from_folder = Folder("/tmp/")
+        to_folder = Folder("s3://mock-bucket/")
+        from_folder = Folder("/tmp/")
 
-    bucket_name = furl(to_folder.uri).origin[5:]
+        bucket_name = furl(to_folder.uri).origin[5:]
 
-    boto3_client_mock().list_objects.return_value = {
-        "Contents": [
-            {"Key": "test.csv"},
-        ]
-    }
+        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"
+        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(
-        "/tmp/mock_join", bucket_name, "mock_join"
-    )
+        callable_func = S3(**self.MOCK_STRATEGY_CONFIG).upload(from_folder, to_folder)
+        callable_func()
+        boto3_client_mock().upload_file.assert_called_once_with(
+            "/tmp/mock_join", bucket_name, "mock_join"
+        )