activeloopai · kristinagrig06 · Sep 13, 2021 · Aug 20, 2021 · Aug 20, 2021 · Aug 20, 2021
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -176,7 +176,7 @@ commands:
                 command: |
                   $Env:GOOGLE_APPLICATION_CREDENTIALS = $Env:CI_GCS_PATH
                   setx /m GOOGLE_APPLICATION_CREDENTIALS "$Env:GOOGLE_APPLICATION_CREDENTIALS"
-                  python3 -m pytest --cov-report=xml --cov=./ --local --s3 --hub-cloud --kaggle --ignore-glob=buH/*
+                  python3 -m pytest --cov-report=xml --cov=./ --local --s3 --gcs --hub-cloud --kaggle --ignore-glob=buH/*
       - when:
           condition: << parameters.unix-like >>
           steps:
@@ -186,9 +186,8 @@ commands:
                   BUGGER_OFF: "true"
                 command: |
                   export GOOGLE_APPLICATION_CREDENTIALS=$HOME/.secrets/gcs.json
-                  python3 -m pytest --cov-report=xml --cov=./ --local --s3 --hub-cloud --kaggle --ignore-glob=buH/*
-
-          parallelism: 10
+                  python3 -m pytest --cov-report=xml --cov=./ --local --s3 --gcs --hub-cloud --kaggle --ignore-glob=buH/*
+
   run-backwards-compatibility-tests:
     steps:
       - run:
@@ -293,6 +292,7 @@ jobs:
           condition: << parameters.unix-like >>
           steps:
             - slack-status
+  no_output_timeout: 30m
   deploy:
     executor: linux
     environment:

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -31,6 +31,7 @@ pip3 install -r hub/requirements/tests.txt
 - `pytest .`: Run all tests with memory only.
 - `pytest . --local`: Run all tests with memory and local.
 - `pytest . --s3`: Run all tests with memory and s3.
+- `pytest . --gcs`: Run all tests with memory and GCS 
 - `pytest . --kaggle`: Run all tests that use the kaggle API.
 - `pytest . --memory-skip --hub-cloud`: Run all tests with hub cloud only.
 #### Backwards Compatibility Tests
@@ -41,6 +42,7 @@ Combine any of the following options to suit your test cases.
 
 - `--local`: Enable local tests.
 - `--s3`: Enable S3 tests.
+- `--gcs`: Enable GCS tests.
 - `--hub-cloud`: Enable hub cloud tests.
 - `--memory-skip`: Disable memory tests.
 - `--s3-path`: Specify an s3 path if you don't have access to our internal testing bucket.

diff --git a/conftest.py b/conftest.py
@@ -28,6 +28,9 @@ def pytest_addoption(parser):
         LOCAL_OPT, action="store_true", help="Local tests will run if enabled."
     )
     parser.addoption(S3_OPT, action="store_true", help="S3 tests will run if enabled.")
+    parser.addoption(
+        GCS_OPT, action="store_true", help="GCS tests will run if enabled."
+    )
     parser.addoption(
         HUB_CLOUD_OPT, action="store_true", help="Hub cloud tests will run if enabled."
     )

diff --git a/hub/constants.py b/hub/constants.py
@@ -88,17 +88,20 @@
 PYTEST_MEMORY_PROVIDER_BASE_ROOT = "mem://hub_pytest"
 PYTEST_LOCAL_PROVIDER_BASE_ROOT = "/tmp/hub_pytest/"  # TODO: may fail for windows
 PYTEST_S3_PROVIDER_BASE_ROOT = "s3://hub-2.0-tests/"
+PYTEST_GCS_PROVIDER_BASE_ROOT = "gcs://snark-test/"
 PYTEST_HUB_CLOUD_PROVIDER_BASE_ROOT = f"hub://{HUB_CLOUD_DEV_USERNAME}/"
 
 # environment variables
 ENV_HUB_DEV_PASSWORD = "ACTIVELOOP_HUB_PASSWORD"
 ENV_KAGGLE_USERNAME = "KAGGLE_USERNAME"
 ENV_KAGGLE_KEY = "KAGGLE_KEY"
+ENV_GOOGLE_APPLICATION_CREDENTIALS = "GOOGLE_APPLICATION_CREDENTIALS"
 
 # pytest options
 MEMORY_OPT = "--memory-skip"
 LOCAL_OPT = "--local"
 S3_OPT = "--s3"
+GCS_OPT = "--gcs"
 HUB_CLOUD_OPT = "--hub-cloud"
 S3_PATH_OPT = "--s3-path"
 KEEP_STORAGE_OPT = "--keep-storage"

diff --git a/hub/core/storage/gcs.py b/hub/core/storage/gcs.py
@@ -0,0 +1,127 @@
+import posixpath
+import json
+import os
+from typing import Dict, Union
+
+from hub.core.storage.provider import StorageProvider
+from google.cloud import storage  # type: ignore
+from google.oauth2 import service_account  # type: ignore
+
+
+class GCSProvider(StorageProvider):
+    """Provider class for using GC storage."""
+
+    def __init__(
+        self,
+        root: str,
+        token: Union[str, Dict] = None,
+    ):
+        """Initializes the GCSProvider
+
+        Example:
+            gcs_provider = GCSProvider("snark-test/gcs_ds")
+
+        Args:
+            root (str): The root of the provider. All read/write request keys will be appended to root.
+            token (str/Dict): GCP token, used for fetching credentials for storage).
+        """
+        self.root = root
+        self.token: Union[str, Dict, None] = token
+        self.missing_exceptions = (
+            FileNotFoundError,
+            IsADirectoryError,
+            NotADirectoryError,
+            AttributeError,
+        )
+        self._initialize_provider()
+
+    def _initialize_provider(self):
+        self._set_bucket_and_path()
+        if self.token:
+            if isinstance(self.token, dict):
+                token_path = posixpath.expanduser("gcs.json")
+                with open(token_path, "wb") as f:
+                    json.dump(self.token, f)
+                self.token = token_path
+            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.token
+            credentials = service_account.Credentials.from_service_account_file(
+                self.token
+            )
+
+            scoped_credentials = credentials.with_scopes(
+                ["https://www.googleapis.com/auth/cloud-platform"]
+            )
+
+            client = storage.Client(credentials=scoped_credentials)
+        else:
+            client = storage.Client(credentials=self.token)
+        self.client_bucket = client.get_bucket(self.bucket)
+
+    def _set_bucket_and_path(self):
+        root = self.root.replace("gcp://", "").replace("gcs://", "")
+        self.bucket = root.split("/")[0]
+        self.path = root
+        if not self.path.endswith("/"):
+            self.path += "/"
+
+    def _get_path_from_key(self, key):
+        return posixpath.join(self.path, key)
+
+    def _list_keys(self):
+        self._blob_objects = self.client_bucket.list_blobs(prefix=self.path)
+        return [obj.name for obj in self._blob_objects]
+
+    def clear(self):
+        """Remove all keys below root - empties out mapping"""
+        self.check_readonly()
+        blob_objects = self.client_bucket.list_blobs(prefix=self.path)
+        for blob in blob_objects:
+            blob.delete()
+
+    def __getitem__(self, key):
+        """Retrieve data"""
+        try:
+            blob = self.client_bucket.get_blob(self._get_path_from_key(key))
+            return blob.download_as_bytes()
+        except self.missing_exceptions:
+            raise KeyError(key)
+
+    def __setitem__(self, key, value):
+        """Store value in key"""
+        self.check_readonly()
+        blob = self.client_bucket.blob(self._get_path_from_key(key))
+        if isinstance(value, memoryview):
+            value = value.tobytes()
+        elif isinstance(value, bytearray):
+            value = bytes(value)
+        blob.upload_from_string(value)
+
+    def __iter__(self):
+        """Iterating over the structure"""
+        yield from [f for f in self._list_keys() if not f.endswith("/")]
+
+    def __len__(self):
+        """Returns length of the structure"""
+        return len(self._list_keys())
+
+    def __delitem__(self, key):
+        """Remove key"""
+        self.check_readonly()
+        blob = self.client_bucket.blob(self._get_path_from_key(key))
+        blob.delete()
+
+    def __contains__(self, key):
+        """Does key exist in mapping?"""
+        stats = storage.Blob(
+            bucket=self.client_bucket, name=self._get_path_from_key(key)
+        ).exists(self.client_bucket.client)
+        return stats
+
+    def __getstate__(self):
+        return (self.root, self.token, self.missing_exceptions)
+
+    def __setstate__(self, state):
+        self.root = state[0]
+        self.token = state[1]
+        self.missing_exceptions = state[2]
+        self._initialize_provider()
diff --git a/hub/requirements/common.txt b/hub/requirements/common.txt
@@ -2,6 +2,8 @@ numpy
 pillow==8.2.0
 boto3
 boto3-stubs[essential]
+google-cloud-storage~=1.42.0
+google-auth~=2.0.1
 pathos
 humbug>=0.2.6
 types-requests

diff --git a/hub/tests/dataset_fixtures.py b/hub/tests/dataset_fixtures.py
@@ -3,13 +3,13 @@
 
 enabled_datasets = pytest.mark.parametrize(
     "ds",
-    ["memory_ds", "local_ds", "s3_ds"],
+    ["memory_ds", "local_ds", "s3_ds", "gcs_ds"],
     indirect=True,
 )
 
 enabled_persistent_dataset_generators = pytest.mark.parametrize(
     "ds_generator",
-    ["local_ds_generator", "s3_ds_generator"],
+    ["local_ds_generator", "s3_ds_generator", "gcs_ds_generator"],
     indirect=True,
 )
 
@@ -45,6 +45,19 @@ def generate_s3_ds():
     return generate_s3_ds
 
 
+@pytest.fixture
+def gcs_ds(gcs_ds_generator):
+    return gcs_ds_generator()
+
+
+@pytest.fixture
+def gcs_ds_generator(gcs_path, gcs_creds):
+    def generate_gcs_ds():
+        return hub.dataset(gcs_path, creds=gcs_creds)
+
+    return generate_gcs_ds
+
+
 @pytest.fixture
 def hub_cloud_ds(hub_cloud_ds_generator):
     return hub_cloud_ds_generator()

diff --git a/hub/tests/path_fixtures.py b/hub/tests/path_fixtures.py
@@ -1,3 +1,4 @@
+from hub.core.storage.gcs import GCSProvider
 from hub.util.storage import storage_provider_from_hub_path
 from hub.core.storage.s3 import S3Provider
 from hub.core.storage.local import LocalProvider
@@ -8,10 +9,13 @@
     KEEP_STORAGE_OPT,
     LOCAL_OPT,
     MEMORY_OPT,
+    PYTEST_GCS_PROVIDER_BASE_ROOT,
     PYTEST_HUB_CLOUD_PROVIDER_BASE_ROOT,
     PYTEST_LOCAL_PROVIDER_BASE_ROOT,
     PYTEST_MEMORY_PROVIDER_BASE_ROOT,
     S3_OPT,
+    GCS_OPT,
+    ENV_GOOGLE_APPLICATION_CREDENTIALS,
 )
 import posixpath
 from hub.tests.common import (
@@ -25,6 +29,7 @@
 MEMORY = "memory"
 LOCAL = "local"
 S3 = "s3"
+GCS = "gcs"
 HUB_CLOUD = "hub_cloud"
 
 
@@ -49,6 +54,12 @@ def _get_path_composition_configs(request):
             "is_id_prefix": True,
             "use_underscores": False,
         },
+        GCS: {
+            "base_root": PYTEST_GCS_PROVIDER_BASE_ROOT,
+            "use_id": True,
+            "is_id_prefix": True,
+            "use_underscores": False,
+        },
         HUB_CLOUD: {
             "base_root": PYTEST_HUB_CLOUD_PROVIDER_BASE_ROOT,
             "use_id": True,
@@ -125,6 +136,27 @@ def s3_path(request):
         S3Provider(path).clear()
 
 
+@pytest.fixture(scope="session")
+def gcs_creds():
+    return os.environ[ENV_GOOGLE_APPLICATION_CREDENTIALS]
+
+
+@pytest.fixture
+def gcs_path(request, gcs_creds):
+    if not is_opt_true(request, GCS_OPT):
+        pytest.skip()
+        return
+
+    path = _get_storage_path(request, GCS)
+    GCSProvider(path, token=gcs_creds).clear()
+
+    yield path
+
+    # clear storage unless flagged otherwise
+    if not is_opt_true(request, KEEP_STORAGE_OPT):
+        GCSProvider(path, token=gcs_creds).clear()
+
+
 @pytest.fixture
 def hub_cloud_path(request, hub_cloud_dev_token):
     if not is_opt_true(request, HUB_CLOUD_OPT):

diff --git a/hub/tests/storage_fixtures.py b/hub/tests/storage_fixtures.py
@@ -1,3 +1,4 @@
+from hub.core.storage.gcs import GCSProvider
 from hub.util.storage import storage_provider_from_hub_path
 from hub.core.storage.s3 import S3Provider
 from hub.core.storage.local import LocalProvider
@@ -7,13 +8,13 @@
 
 enabled_storages = pytest.mark.parametrize(
     "storage",
-    ["memory_storage", "local_storage", "s3_storage"],
+    ["memory_storage", "local_storage", "s3_storage", "gcs_storage"],
     indirect=True,
 )
 
 enabled_persistent_storages = pytest.mark.parametrize(
     "storage",
-    ["local_storage", "s3_storage"],
+    ["local_storage", "s3_storage", "gcs_storage"],
     indirect=True,
 )
 
@@ -33,6 +34,11 @@ def s3_storage(s3_path):
     return S3Provider(s3_path)
 
 
+@pytest.fixture
+def gcs_storage(gcs_path):
+    return GCSProvider(gcs_path)
+
+
 @pytest.fixture
 def hub_cloud_storage(hub_cloud_path, hub_cloud_dev_token):
     return storage_provider_from_hub_path(hub_cloud_path, token=hub_cloud_dev_token)

diff --git a/hub/util/keys.py b/hub/util/keys.py
@@ -42,7 +42,7 @@ def dataset_exists(storage: StorageProvider) -> bool:
     try:
         storage[get_dataset_meta_key()]
         return True
-    except KeyError:
+    except (KeyError, FileNotFoundError):
         return False
 
 

diff --git a/hub/util/storage.py b/hub/util/storage.py
@@ -1,3 +1,4 @@
+from hub.core.storage.gcs import GCSProvider
 from hub.util.cache_chain import generate_chain
 from hub.constants import MB
 from hub.util.tag import check_hub_path
@@ -25,6 +26,7 @@ def storage_provider_from_path(
 
     Returns:
         If given a path starting with s3://  returns the S3Provider.
+        If given a path starting with gcp:// or gcs:// returns the GCPProvider.
         If given a path starting with mem:// returns the MemoryProvider.
         If given a path starting with hub:// returns the underlying cloud Provider.
         If given a valid local path, returns the LocalProvider.
@@ -43,6 +45,8 @@ def storage_provider_from_path(
         storage: StorageProvider = S3Provider(
             path, key, secret, session_token, endpoint_url, region, token=token
         )
+    elif path.startswith("gcp://") or path.startswith("gcs://"):
+        storage = GCSProvider(path, creds)
     elif path.startswith("mem://"):
         storage = MemoryProvider(path)
     elif path.startswith("hub://"):