[Python] Catch PermissionDenied when checking the feature store in VertexAIFeatureStoreEnrichmentHandler

CHANGES.md

@@ -87,6 +87,7 @@
 * (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)).
 * (Java) Fixed the declared schema of the error output of the Kafka write SchemaTransform, which wrapped the error schema a second time and did not match the rows it emits ([#39760](https://github.com/apache/beam/issues/39760)).
 * (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed in recent releases, which broke builds of Go modules depending on a current `genproto` version ([#40018](https://github.com/apache/beam/issues/40018)).
+* (Python) `VertexAIFeatureStoreEnrichmentHandler` now logs a warning instead of failing when the caller lacks permission to read the feature store metadata ([#40442](https://github.com/apache/beam/issues/40442)).
 * (Java) BigQueryIO now treats a 404 when deleting a temporary table or dataset as success, so a replayed work item whose earlier attempt already deleted it no longer retries forever ([#24997](https://github.com/apache/beam/issues/24997)).
 * (Java) IcebergIO now writes rows containing `EnumerationType` (proto enum) fields as strings, instead of throwing `Unsupported Beam logical type Enum` ([#40299](https://github.com/apache/beam/issues/40299)).
 * (Python) Fixed stateful DoFns with side inputs sometimes taking the timer key coder from a side input instead of the main input, which could make the worker fail to decode timer keys with `Unknown type tag` ([#40374](https://github.com/apache/beam/issues/40374)).

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

@@ -18,6 +18,7 @@
 
 import proto
 from google.api_core.exceptions import NotFound
+from google.api_core.exceptions import PermissionDenied
 from google.cloud import aiplatform
 
 import apache_beam as beam
@@ -113,23 +114,23 @@ 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(
           project=self.project,
           location=self.location,
           feature_online_store=self.feature_store_name)
+      # creating the client doesn't check permissions, this RPC does.
       feature_store = admin_client.get_feature_online_store(
           name=feature_store_path)
-
+    except PermissionDenied:
+      _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:
       if not feature_store:
         raise NotFound(
             'Vertex AI Feature Store %s does not exists in %s' %

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

@@ -15,8 +15,12 @@
 # limitations under the License.
 #
 import unittest
+from unittest import mock
 
 try:
+  from google.api_core.exceptions import NotFound
+  from google.api_core.exceptions import PermissionDenied
+
   from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import VertexAIFeatureStoreEnrichmentHandler
   from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import \
       VertexAIFeatureStoreLegacyEnrichmentHandler
@@ -39,6 +43,42 @@ def test_raise_error_duplicate_api_endpoint_online_store(self):
           client_options={'api_endpoint': 'region@google.com'},
       )
 
+  @mock.patch(
+      'apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.'
+      'aiplatform.gapic.FeatureOnlineStoreAdminServiceClient')
+  def test_permission_denied_online_store(self, mock_admin_client):
+    mock_admin_client.return_value.get_feature_online_store.side_effect = (
+        PermissionDenied('permission denied'))
+    with self.assertLogs(
+        'apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store',
+        level='WARNING') as logs:
+      _ = 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.assertIn('insufficient admin permission', logs.output[0])
+
+  @mock.patch(
+      'apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.'
+      'aiplatform.gapic.FeatureOnlineStoreAdminServiceClient')
+  def test_raise_error_feature_store_not_found_online_store(
+      self, mock_admin_client):
+    mock_admin_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',
+      )
+
   def test_raise_error_duplicate_api_endpoint_legacy_store(self):
     with self.assertRaises(ValueError):
       _ = VertexAIFeatureStoreLegacyEnrichmentHandler(