Repository: AgnostiqHQ/covalent
Pull request #1699 (2023-06-12 19:33:33)

Add google storage file transfer strategy
<!--
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Add a note to /CHANGELOG.md summarizing the changes.
⚠️ If your pull request fixes an open issue, please link to the issue.
⚠️ Ensure that your branch is up-to-date with the base branch.
-->
This pull request adds file transfer support for Google Cloud storage.  To use this users will need to have set up a bucket in a Google Cloud storage account (with a billing account attached).  Additionally users will need to have created a service account + keys for machine-based authentication.  During testing I used the `Storage Admin` role for the service account.  

Syntax is the following:

```
import covalent as ct

gcloud_strategy = ct.fs_strategies.GCloud(
    credentials="/path/to/service_account_credentials.json",
    project_id="my_gcp_project_id",
)

input_file = ct.fs.TransferFromRemote(
    "gs://<remote-bucket>/<filename>",
    "/local/path/to/downloaded/object",
    strategy=gcloud_strategy,
)
output_file = ct.fs.TransferToRemote(
    "gs://<remote-bucket>/<uploaded_filename>",
    "/local/path/to/object",
    strategy=gcloud_strategy,
)

@ct.electron(files=[input_file, output_file])
def gcloud_transfer(files):
    file_path = files[0][1]
    return os.path.exists(file_path)

@ct.lattice
def test_gcloud_transfers():
    return gcloud_transfer()

dispatch_id = ct.dispatch(test_gcloud_transfers)()
print(dispatch_id)
```

Note that we assume URIs in Google Cloud Storage use the format `gs://<bucket_name>/<blob_name>` and recognize `gs` as the file scheme of remote objects.

- [x] I have added the tests to cover my changes.
- [x] I have updated the documentation and CHANGELOG accordingly.
- [x] I have read the CONTRIBUTING document.


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1699?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) Report
> Merging [#1699](https://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1699?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (ba998cb) into [develop](https://app.codecov.io/gh/AgnostiqHQ/covalent/commit/935553d6e79f1919e239e4441d5be17d60ce1d34?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (935553d) will **decrease** coverage by `0.24%`.
> The diff coverage is `46.15%`.

> :exclamation: Current head ba998cb differs from pull request most recent head 0a06283. Consider uploading reports for the commit 0a06283 to get more accurate results

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


[![Impacted file tree graph](https://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1699/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://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1699?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)

```diff
@@             Coverage Diff             @@
##           develop    #1699      +/-   ##
===========================================
- Coverage    87.95%   87.72%   -0.24%     
===========================================
  Files          160      161       +1     
  Lines         6876     6915      +39     
  Branches       151      151              
===========================================
+ Hits          6048     6066      +18     
- Misses         723      744      +21     
  Partials       105      105              
```

| Flag | Coverage Δ | |
|---|---|---|
| Dispatcher | `87.57% <ø> (ø)` | |
| Functional_Tests | `58.09% <46.15%> (-0.10%)` | :arrow_down: |
| SDK | `90.80% <46.15%> (-0.54%)` | :arrow_down: |
| UI_Frontend | `72.64% <ø> (ø)` | |

</details>

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fe718f95..f10fd38b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- File transfer strategy for GCP storage
 - Add CLI status for zombie, stopped process.
 
 ### Docs
@@ -39,7 +40,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix `Lattice` docstring in RTD.
 - Added a missing `,` to the Slurm docs.
 
-
 ### Operations
 
 - Update the PR template.
diff --git a/covalent/_file_transfer/enums.py b/covalent/_file_transfer/enums.py
index a1ba9a51..6353da08 100644
--- a/covalent/_file_transfer/enums.py
+++ b/covalent/_file_transfer/enums.py
@@ -30,6 +30,7 @@ class FileSchemes(str, enum.Enum):
     File = "file"
     S3 = "s3"
     Blob = "https"
+    GCloud = "gs"
     Globus = "globus"
     HTTP = "http"
     HTTPS = "https"
@@ -40,6 +41,7 @@ class FileTransferStrategyTypes(str, enum.Enum):
     Rsync = "rsync"
     HTTP = "http"
     S3 = "s3"
+    GCloud = "gs"
     FTP = "ftp"
     GLOBUS = "globus"
 
@@ -49,6 +51,7 @@ SchemeToStrategyMap = {
     "http": FileTransferStrategyTypes.HTTP,
     "https": FileTransferStrategyTypes.HTTP,
     "s3": FileTransferStrategyTypes.S3,
+    "gs": FileTransferStrategyTypes.GCloud,
     "ftp": FileTransferStrategyTypes.FTP,
     "globus": FileTransferStrategyTypes.GLOBUS,
 }
diff --git a/covalent/_file_transfer/file.py b/covalent/_file_transfer/file.py
index 8faea355..5a867e2e 100644
--- a/covalent/_file_transfer/file.py
+++ b/covalent/_file_transfer/file.py
@@ -87,6 +87,7 @@ class File:
         return self._is_remote or self.scheme in [
             FileSchemes.S3,
             FileSchemes.Blob,
+            FileSchemes.GCloud,
             FileSchemes.Globus,
             FileSchemes.HTTP,
             FileSchemes.HTTPS,
@@ -139,6 +140,8 @@ class File:
             return FileSchemes.S3
         if scheme == FileSchemes.Blob and "blob.core.windows.net" in host:
             return FileSchemes.Blob
+        if scheme == FileSchemes.GCloud:
+            return FileSchemes.GCloud
         if scheme == FileSchemes.FTP:
             return FileSchemes.FTP
         if scheme == FileSchemes.HTTP:
diff --git a/covalent/_file_transfer/strategies/__init__.py b/covalent/_file_transfer/strategies/__init__.py
index e0ae22f2..c2a7ba0a 100644
--- a/covalent/_file_transfer/strategies/__init__.py
+++ b/covalent/_file_transfer/strategies/__init__.py
@@ -30,3 +30,8 @@ try:
     from .blob_strategy import Blob
 except:  # pragma: no cover
     pass
+
+try:
+    from .gcloud_strategy import GCloud
+except:  # pragma: no cover
+    pass
diff --git a/covalent/_file_transfer/strategies/gcloud_strategy.py b/covalent/_file_transfer/strategies/gcloud_strategy.py
new file mode 100644
index 00000000..f98adf7c
--- /dev/null
+++ b/covalent/_file_transfer/strategies/gcloud_strategy.py
@@ -0,0 +1,172 @@
+# 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 json
+from pathlib import Path
+from typing import Callable, Tuple
+
+from furl import furl
+
+from ..._shared_files import logger
+from .. import File
+from .transfer_strategy_base import FileTransferStrategy
+
+app_log = logger.app_log
+
+
+class GCloud(FileTransferStrategy):
+    """Implements FileTransferStrategy class to transfer files to/from Google Cloud Storage.
+
+    Args:
+        credentials: Path to OAuth 2.0 credentials JSON file for a service account
+        project_id: ID of a project in GCP
+
+    Attributes:
+        credentials: String containing OAuth 2.0 credentials
+        project_id: ID of a project in GCP
+    """
+
+    def __init__(self, credentials: str = None, project_id: str = None):
+        if credentials is not None:
+            credentials_json = Path(credentials).resolve()
+
+            if not credentials_json.is_file():
+                raise ValueError("Cloud not locate credentials for GCloud file transfers.")
+
+            with open(credentials_json) as f:
+                self.credentials = json.load(f)
+        else:
+            self.credentials = None
+
+        self.project_id = project_id
+
+    def _get_service_client(self, bucket_name: str):
+        """Returns the service client object for the Blob storage account.
+
+        Args:
+            bucket_name: Name of a storage bucket in Google Cloud
+
+        Returns:
+            Storage client object for the Google Cloud storage service
+        """
+
+        try:
+            from google.cloud import storage
+            from google.oauth2 import service_account
+        except ImportError:
+            print("Missing required dependency. Reinstall package as `pip install covalent[gcp]`.")
+            raise
+
+        return storage.Client(
+            project=self.project_id,
+            credentials=service_account.Credentials.from_service_account_info(self.credentials),
+        ).bucket(bucket_name)
+
+    def _parse_uri(self, gcloud_uri: str) -> Tuple[str, str]:
+        """Parses a GCloud object URI and returns the bucket name and the object path.
+
+        Args:
+            blob_uri: A URI for an Azure Blob object in the form https://<storage_account_name>.blob.core.windows.net/<container_name>/<blob_name>
+
+        Returns:
+            parsed_uri: A tuple containing the storage account name, container name, and blob name
+        """
+
+        object_furl = furl(gcloud_uri)
+
+        bucket_name = object_furl.host
+        object_path = object_furl.path
+
+        if str(object_path).startswith("/"):
+            object_path = Path(str(object_path)[1:])
+
+        return (bucket_name, object_path)
+
+    def _download_file(self, blob_client, destination_path: str) -> None:
+        blob_client.download_to_filename(destination_path)
+
+    def download(self, from_file: File, to_file: File = File()) -> Callable:
+        from_filepath = from_file.filepath
+        to_filepath = to_file.filepath
+
+        bucket_name, object_path = self._parse_uri(from_file.uri)
+
+        app_log.debug(
+            f"GCloud download; bucket name: {bucket_name}, from_filepath: {from_filepath}, to_filepath {to_filepath}."
+        )
+
+        def callable():
+            """Download file or directory from a Google Cloud Storage bucket."""
+            from pathlib import Path
+
+            service_client = self._get_service_client(bucket_name)
+
+            blobs = service_client.list_blobs(prefix=object_path)
+            for blob in blobs:
+                if blob.name.endswith("/"):
+                    continue
+
+                dest_obj_path = (
+                    Path(to_filepath) / Path(blob.name).relative_to(object_path)
+                    if to_file.is_dir
+                    else Path(to_filepath)
+                )
+                dest_obj_path.parent.mkdir(parents=True, exist_ok=True)
+                self._download_file(blob, str(dest_obj_path))
+
+        return callable
+
+    def _upload_file(self, service_client, source_path: str, destination_path: Path) -> None:
+        blob = service_client.blob(destination_path)
+        blob.upload_from_filename(source_path)
+
+    def upload(self, from_file: File, to_file: File = File()) -> Callable:
+        from_filepath = from_file.filepath
+        to_filepath = to_file.filepath
+
+        bucket_name, object_path = self._parse_uri(to_file.uri)
+
+        app_log.debug(
+            f"GCloud upload; bucket name: {bucket_name}, from_filepath: {from_filepath}, to_filepath {to_filepath}."
+        )
+
+        def callable():
+            """Upload file or directory to a Google Cloud Storage bucket."""
+            from pathlib import Path
+
+            service_client = self._get_service_client(bucket_name)
+
+            files = (
+                [path for path in Path(from_filepath).rglob("*") if path.is_file()]
+                if from_file.is_dir
+                else [from_filepath]
+            )
+            for file in files:
+                dest_obj_path = (
+                    Path(object_path) / Path(file).relative_to(from_filepath)
+                    if to_file.is_dir
+                    else Path(object_path)
+                )
+                self._upload_file(service_client, file, str(dest_obj_path))
+
+        return callable
+
+    def cp(self, from_file: File, to_file: File = File()) -> File:  # pragma: no cover
+        raise NotImplementedError
diff --git a/setup.py b/setup.py
index 2242514a..adc9426a 100644
--- a/setup.py
+++ b/setup.py
@@ -214,6 +214,7 @@ setup_info = {
     "extras_require": {
         "aws": ["boto3>=1.20.48"],
         "azure": ["azure-identity>=1.13.0", "azure-storage-blob>=12.16.0"],
+        "gcp": ["google-auth>=2.16.2", "google-cloud-storage>=2.7.0"],
         "mysql": ["mysqlclient>=2.1.1"],
         "postgres": ["psycopg2-binary>=2.9.5"],
     },


Test patch:
diff --git a/tests/covalent_tests/file_transfer/strategies/blob_strategy_test.py b/tests/covalent_tests/file_transfer/strategies/blob_strategy_test.py
index 1459ea69..93ac117b 100644
--- a/tests/covalent_tests/file_transfer/strategies/blob_strategy_test.py
+++ b/tests/covalent_tests/file_transfer/strategies/blob_strategy_test.py
@@ -82,7 +82,6 @@ def test_get_blob_client(mocker, blob_strategy):
 
 def test_parse_blob_uri(blob_strategy):
     result = blob_strategy._parse_blob_uri(MOCK_REMOTE_FILEPATH)
-    print(result)
 
     assert result[0] == MOCK_BLOB_STORAGE_ACCOUNT_URL
     assert result[1] == MOCK_BLOB_CONTAINER
diff --git a/tests/covalent_tests/file_transfer/strategies/gcloud_strategy_test.py b/tests/covalent_tests/file_transfer/strategies/gcloud_strategy_test.py
new file mode 100644
index 00000000..7d14347b
--- /dev/null
+++ b/tests/covalent_tests/file_transfer/strategies/gcloud_strategy_test.py
@@ -0,0 +1,192 @@
+# 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 sys
+from pathlib import Path
+from unittest.mock import MagicMock, mock_open
+
+import pytest
+
+from covalent._file_transfer import File
+from covalent._file_transfer.strategies.gcloud_strategy import GCloud
+
+MOCK_LOCAL_FILEPATH = "/Users/user/data.csv"
+MOCK_STORAGE_BUCKET = "mock-bucket"
+MOCK_REMOTE_OBJECT_NAME = "covalent-tmp/data.csv"
+MOCK_REMOTE_FILEPATH = f"gs://{MOCK_STORAGE_BUCKET}/{MOCK_REMOTE_OBJECT_NAME}"
+MOCK_CREDENTIALS_PATH = "/path/to/mock/credentials"
+MOCK_PROJECT_ID = "mock-project-id"
+MOCK_CREDENTIALS = {
+    "mock-key": "mock-value",
+}
+
+
+@pytest.fixture
+def gcloud_strategy():
+    strategy = GCloud()
+    strategy.credentials = MOCK_CREDENTIALS
+    strategy.project_id = MOCK_PROJECT_ID
+
+    return strategy
+
+
+def test_init(mocker):
+    mocker.patch("pathlib.Path.is_file", return_value=True)
+    open_mock = mocker.patch(
+        "covalent._file_transfer.strategies.gcloud_strategy.open", mock_open()
+    )
+    json_load_mock = mocker.patch("json.load", return_value=MOCK_CREDENTIALS)
+
+    strategy = GCloud(MOCK_CREDENTIALS_PATH, MOCK_PROJECT_ID)
+
+    open_mock.assert_called_once_with(Path(MOCK_CREDENTIALS_PATH).resolve())
+    json_load_mock.assert_called_once_with(open_mock())
+
+    assert strategy.credentials == MOCK_CREDENTIALS
+    assert strategy.project_id == MOCK_PROJECT_ID
+
+    strategy = GCloud()
+
+    assert strategy.credentials is None
+    assert strategy.project_id is None
+
+    mocker.patch("pathlib.Path.is_file", return_value=False)
+    with pytest.raises(ValueError):
+        strategy = GCloud(MOCK_CREDENTIALS_PATH, MOCK_PROJECT_ID)
+
+
+def test_get_service_client(mocker, gcloud_strategy):
+    sys.modules["google"] = MagicMock()
+    cloud_mock = MagicMock()
+    sys.modules["google.cloud"] = cloud_mock
+    oauth2_mock = MagicMock()
+    sys.modules["google.oauth2"] = oauth2_mock
+
+    storage_client_mock = cloud_mock.storage.Client
+    service_account_mock = oauth2_mock.service_account.Credentials.from_service_account_info
+    bucket_mock = storage_client_mock().bucket
+
+    client = gcloud_strategy._get_service_client(MOCK_STORAGE_BUCKET)
+
+    service_account_mock.assert_called_once_with(MOCK_CREDENTIALS)
+    storage_client_mock.assert_called_with(
+        project=MOCK_PROJECT_ID, credentials=service_account_mock.return_value
+    )
+    bucket_mock.assert_called_once_with(MOCK_STORAGE_BUCKET)
+
+    assert client == bucket_mock.return_value
+
+    sys.modules["google.oauth2"] = None
+    with pytest.raises(ImportError):
+        gcloud_strategy._get_service_client(MOCK_STORAGE_BUCKET)
+
+    sys.modules["google.cloud"] = None
+    with pytest.raises(ImportError):
+        gcloud_strategy._get_service_client(MOCK_STORAGE_BUCKET)
+
+
+def test_parse_uri(gcloud_strategy):
+    result = gcloud_strategy._parse_uri(MOCK_REMOTE_FILEPATH)
+
+    assert result[0] == MOCK_STORAGE_BUCKET
+    assert str(result[1]) == MOCK_REMOTE_OBJECT_NAME
+
+
+def test_download_file(gcloud_strategy):
+    blob_client_mock = MagicMock()
+    destination_path = MOCK_LOCAL_FILEPATH
+    download_mock = blob_client_mock.download_to_filename
+
+    gcloud_strategy._download_file(blob_client_mock, destination_path)
+
+    download_mock.assert_called_once_with(destination_path)
+
+
+def test_download(mocker, gcloud_strategy):
+    from_file = File(MOCK_REMOTE_FILEPATH)
+    to_file = File(MOCK_LOCAL_FILEPATH)
+
+    mocker.patch("covalent._file_transfer.strategies.gcloud_strategy.app_log")
+
+    parse_uri_mock = mocker.patch(
+        "covalent._file_transfer.strategies.gcloud_strategy.GCloud._parse_uri"
+    )
+    parse_uri_mock.return_value = (MOCK_STORAGE_BUCKET, MOCK_REMOTE_OBJECT_NAME)
+
+    get_service_client_mock = mocker.patch(
+        "covalent._file_transfer.strategies.gcloud_strategy.GCloud._get_service_client"
+    )
+    list_blobs_mock = get_service_client_mock().list_blobs
+    blob_mock = MagicMock()
+    list_blobs_mock.return_value = [blob_mock]
+    blob_mock.name = MOCK_REMOTE_OBJECT_NAME
+
+    mkdir_mock = mocker.patch("pathlib.Path.mkdir")
+    download_file_mock = mocker.patch(
+        "covalent._file_transfer.strategies.gcloud_strategy.GCloud._download_file"
+    )
+
+    gcloud_strategy.download(from_file, to_file)()
+
+    parse_uri_mock.assert_called_once_with(MOCK_REMOTE_FILEPATH)
+    get_service_client_mock.assert_called_with(MOCK_STORAGE_BUCKET)
+    list_blobs_mock.assert_called_once_with(prefix=MOCK_REMOTE_OBJECT_NAME)
+    mkdir_mock.assert_called_once_with(parents=True, exist_ok=True)
+    download_file_mock.assert_called_once_with(blob_mock, MOCK_LOCAL_FILEPATH)
+
+
+def test_upload_file(gcloud_strategy):
+    service_client = MagicMock()
+    source_path = MOCK_LOCAL_FILEPATH
+    destination_path = MOCK_REMOTE_OBJECT_NAME
+
+    blob_mock = service_client.blob
+    upload_mock = blob_mock().upload_from_filename
+
+    gcloud_strategy._upload_file(service_client, source_path, destination_path)
+
+    blob_mock.assert_called_with(destination_path)
+    upload_mock.assert_called_once_with(source_path)
+
+
+def test_upload(mocker, gcloud_strategy):
+    from_file = File(MOCK_LOCAL_FILEPATH)
+    to_file = File(MOCK_REMOTE_FILEPATH)
+
+    parse_uri_mock = mocker.patch(
+        "covalent._file_transfer.strategies.gcloud_strategy.GCloud._parse_uri"
+    )
+    parse_uri_mock.return_value = (MOCK_STORAGE_BUCKET, MOCK_REMOTE_OBJECT_NAME)
+
+    mocker.patch("covalent._file_transfer.strategies.gcloud_strategy.app_log")
+    get_service_client_mock = mocker.patch(
+        "covalent._file_transfer.strategies.gcloud_strategy.GCloud._get_service_client"
+    )
+    upload_file_mock = mocker.patch(
+        "covalent._file_transfer.strategies.gcloud_strategy.GCloud._upload_file"
+    )
+
+    gcloud_strategy.upload(from_file, to_file)()
+
+    parse_uri_mock.assert_called_once_with(MOCK_REMOTE_FILEPATH)
+    get_service_client_mock.assert_called_with(MOCK_STORAGE_BUCKET)
+    upload_file_mock.assert_called_once_with(
+        get_service_client_mock.return_value, MOCK_LOCAL_FILEPATH, MOCK_REMOTE_OBJECT_NAME
+    )