Improve GCP exception handling (#1561)

woop · web-flow · commit 73ac3de613a7 · 2021-05-15T14:22:52.000-07:00
* Remove try/catch that fails silently

Signed-off-by: Willem Pienaar &lt;git@willem.co&gt;

* Remove unused import

Signed-off-by: Willem Pienaar &lt;git@willem.co&gt;

* Provided clearer exception handling for GCP dependencies

Signed-off-by: Willem Pienaar &lt;git@willem.co&gt;

* Ensure that GCP dependencies aren't loaded unnecessarily

Signed-off-by: Willem Pienaar &lt;git@willem.co&gt;

* Fix lint

Signed-off-by: Willem Pienaar &lt;git@willem.co&gt;
diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py
@@ -1,3 +1,6 @@
+from colorama import Fore, Style
+
+
 class FeastObjectNotFoundException(Exception):
     pass
 
@@ -47,3 +50,13 @@ def __init__(self, module_name, class_name):
         super().__init__(
             f"Could not import provider '{class_name}' from module '{module_name}'"
         )
+
+
+class FeastExtrasDependencyImportError(Exception):
+    def __init__(self, extras_type: str, nested_error: str):
+        message = (
+            nested_error
+            + "\n"
+            + f"You may need run {Style.BRIGHT + Fore.GREEN}pip install 'feast[{extras_type}]'{Style.RESET_ALL}"
+        )
+        super().__init__(message)
diff --git a/sdk/python/feast/infra/gcp.py b/sdk/python/feast/infra/gcp.py
@@ -6,7 +6,6 @@
 import mmh3
 import pandas
 import pyarrow
-from google.auth.exceptions import DefaultCredentialsError
 from tqdm import tqdm
 
 from feast import FeatureTable, utils
@@ -27,6 +26,14 @@
 from feast.registry import Registry
 from feast.repo_config import DatastoreOnlineStoreConfig, RepoConfig
 
+try:
+    from google.auth.exceptions import DefaultCredentialsError
+    from google.cloud import bigquery, datastore
+except ImportError as e:
+    from feast.errors import FeastExtrasDependencyImportError
+
+    raise FeastExtrasDependencyImportError("gcp", str(e))
+
 
 class GcpProvider(Provider):
     _gcp_project_id: Optional[str]
@@ -39,8 +46,6 @@ def __init__(self, config: RepoConfig):
             self._gcp_project_id = None
 
     def _initialize_client(self):
-        from google.cloud import datastore
-
         try:
             if self._gcp_project_id is not None:
                 return datastore.Client(self._gcp_project_id)
@@ -49,7 +54,8 @@ def _initialize_client(self):
         except DefaultCredentialsError as e:
             raise FeastProviderLoginError(
                 str(e)
-                + '\nIt may be necessary to run "gcloud auth application-default login" if you would like to use your local Google Cloud account'
+                + '\nIt may be necessary to run "gcloud auth application-default login" if you would like to use your '
+                "local Google Cloud account "
             )
 
     def update_infra(
@@ -61,7 +67,6 @@ def update_infra(
         entities_to_keep: Sequence[Entity],
         partial: bool,
     ):
-        from google.cloud import datastore
 
         client = self._initialize_client()
 
@@ -190,8 +195,6 @@ def materialize_single_feature_view(
 
     @staticmethod
     def _pull_query(query: str) -> pyarrow.Table:
-        from google.cloud import bigquery
-
         client = bigquery.Client()
         query_job = client.query(query)
         return query_job.to_arrow()
@@ -248,8 +251,6 @@ def _write_minibatch(
     ],
     progress: Optional[Callable[[int], Any]],
 ):
-    from google.cloud import datastore
-
     entities = []
     for entity_key, features, timestamp, created_ts in data:
         document_id = compute_datastore_entity_id(entity_key)
diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -5,7 +5,6 @@
 
 import pandas
 import pyarrow
-from google.auth.exceptions import DefaultCredentialsError
 from jinja2 import BaseLoader, Environment
 
 from feast.data_source import BigQuerySource, DataSource
@@ -20,6 +19,15 @@
 from feast.registry import Registry
 from feast.repo_config import RepoConfig
 
+try:
+    from google.auth.exceptions import DefaultCredentialsError
+    from google.cloud import bigquery
+
+except ImportError as e:
+    from feast.errors import FeastExtrasDependencyImportError
+
+    raise FeastExtrasDependencyImportError("gcp", str(e))
+
 
 class BigQueryOfflineStore(OfflineStore):
     @staticmethod
@@ -192,7 +200,6 @@ class FeatureViewQueryContext:
 
 def _upload_entity_df_into_bigquery(project, entity_df, client) -> str:
     """Uploads a Pandas entity dataframe into a BigQuery table and returns a reference to the resulting table"""
-    from google.cloud import bigquery
 
     # First create the BigQuery dataset if it doesn't exist
     dataset = bigquery.Dataset(f"{client.project}.feast_{project}")
@@ -303,8 +310,6 @@ def build_point_in_time_query(
 
 def _get_bigquery_client():
     try:
-        from google.cloud import bigquery
-
         client = bigquery.Client()
     except DefaultCredentialsError as e:
         raise FeastProviderLoginError(
diff --git a/sdk/python/feast/infra/offline_stores/helpers.py b/sdk/python/feast/infra/offline_stores/helpers.py
@@ -1,8 +1,6 @@
 from typing import List
 
 from feast.data_source import BigQuerySource, DataSource, FileSource
-from feast.infra.offline_stores.bigquery import BigQueryOfflineStore
-from feast.infra.offline_stores.file import FileOfflineStore
 from feast.infra.offline_stores.offline_store import OfflineStore
 
 
@@ -13,10 +11,14 @@ def get_offline_store_from_sources(sources: List[DataSource]) -> OfflineStore:
 
     # Retrieve features from ParquetOfflineStore
     if all(source == FileSource for source in source_types):
+        from feast.infra.offline_stores.file import FileOfflineStore
+
         return FileOfflineStore()
 
     # Retrieve features from BigQueryOfflineStore
     if all(source == BigQuerySource for source in source_types):
+        from feast.infra.offline_stores.bigquery import BigQueryOfflineStore
+
         return BigQueryOfflineStore()
 
     # Could not map inputs to an OfflineStore implementation
diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py
@@ -431,18 +431,13 @@ def _write_registry(self, registry_proto: RegistryProto):
 class GCSRegistryStore(RegistryStore):
     def __init__(self, uri: str):
         try:
-            from google.auth.exceptions import DefaultCredentialsError
             from google.cloud import storage
-        except ImportError:
-            # TODO: Ensure versioning depends on requirements.txt/setup.py and isn't hardcoded
-            raise ImportError(
-                "Install package google-cloud-storage==1.20.* for gcs support"
-                "run ```pip install google-cloud-storage==1.20.*```"
-            )
-        try:
-            self.gcs_client = storage.Client()
-        except DefaultCredentialsError:
-            self.gcs_client = storage.Client.create_anonymous_client()
+        except ImportError as e:
+            from feast.errors import FeastExtrasDependencyImportError
+
+            raise FeastExtrasDependencyImportError("gcp", str(e))
+
+        self.gcs_client = storage.Client()
         self._uri = urlparse(uri)
         self._bucket = self._uri.hostname
         self._blob = self._uri.path.lstrip("/")
