Fix uncaught permission exception in VertexAIFeatureStoreEnrichmentHandler

sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store.py

@@ -113,14 +113,6 @@ def __init__(
     try:
       admin_client = aiplatform.gapic.FeatureOnlineStoreAdminServiceClient(
           **self.kwargs)
-    except Exception:
-      _LOGGER.warning(
-          'Due to insufficient admin permission, could not verify '
-          'the existence of feature store. If the `exception_level` '
-          'is set to WARN then make sure the feature store exists '
-          'otherwise the data enrichment will not happen without '
-          'throwing an error.')
-    else:
       location_path = admin_client.common_location_path(
           project=self.project, location=self.location)
       feature_store_path = admin_client.feature_online_store_path(
@@ -129,11 +121,21 @@ def __init__(
           feature_online_store=self.feature_store_name)
       feature_store = admin_client.get_feature_online_store(
           name=feature_store_path)
+    except NotFound:
+      raise
+    except Exception:
+      _LOGGER.warning(
+          'Due to insufficient admin permission, could not verify '
+          'the existence of feature store. If the `exception_level` '
+          'is set to WARN then make sure the feature store exists '
+          'otherwise the data enrichment will not happen without '
+          'throwing an error.')
+      return
 
-      if not feature_store:
-        raise NotFound(
-            'Vertex AI Feature Store %s does not exists in %s' %
-            (self.feature_store_name, location_path))
+    if not feature_store:
+      raise NotFound(
+          'Vertex AI Feature Store %s does not exists in %s' %
+          (self.feature_store_name, location_path))
 
   def __enter__(self):
     """Connect with the Vertex AI Feature Store."""

sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_test.py

@@ -15,6 +15,8 @@
 # limitations under the License.
 #
 import unittest
+import mock
+from google.api_core.exceptions import NotFound
 
 try:
   from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import VertexAIFeatureStoreEnrichmentHandler
@@ -39,6 +41,50 @@ def test_raise_error_duplicate_api_endpoint_online_store(self):
           client_options={'api_endpoint': 'region@google.com'},
       )
 
+  @mock.patch(
+      'google.cloud.aiplatform.gapic.FeatureOnlineStoreAdminServiceClient')
+  def test_feature_store_not_found(self, mock_client):
+    mock_client.return_value.get_feature_online_store.side_effect = NotFound(
+        'not found')
+    with self.assertRaises(NotFound):
+      _ = VertexAIFeatureStoreEnrichmentHandler(
+          project='project',
+          location='location',
+          api_endpoint='location@google.com',
+          feature_store_name='feature_store',
+          feature_view_name='feature_view',
+          row_key='row_key',
+      )
+
+  @mock.patch(
+      'google.cloud.aiplatform.gapic.FeatureOnlineStoreAdminServiceClient')
+  def test_feature_store_falsy(self, mock_client):
+    mock_client.return_value.get_feature_online_store.return_value = None
+    with self.assertRaises(NotFound):
+      _ = VertexAIFeatureStoreEnrichmentHandler(
+          project='project',
+          location='location',
+          api_endpoint='location@google.com',
+          feature_store_name='feature_store',
+          feature_view_name='feature_view',
+          row_key='row_key',
+      )
+
+  @mock.patch(
+      'google.cloud.aiplatform.gapic.FeatureOnlineStoreAdminServiceClient')
+  def test_feature_store_exception(self, mock_client):
+    mock_client.return_value.get_feature_online_store.side_effect = Exception(
+        'test exception')
+    handler = VertexAIFeatureStoreEnrichmentHandler(
+        project='project',
+        location='location',
+        api_endpoint='location@google.com',
+        feature_store_name='feature_store',
+        feature_view_name='feature_view',
+        row_key='row_key',
+    )
+    self.assertIsNotNone(handler)
+
   def test_raise_error_duplicate_api_endpoint_legacy_store(self):
     with self.assertRaises(ValueError):
       _ = VertexAIFeatureStoreLegacyEnrichmentHandler(