diff --git a/CHANGES.md b/CHANGES.md index 4682f88612d1..eb58a2a85990 100644 --- a/CHANGES.md +++ b/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)). diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store.py b/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store.py index c364378ea8a5..88d2f3cc6ae2 100644 --- a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store.py +++ b/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' % diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_test.py b/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_test.py index 211529be4dc9..233f0e11da6d 100644 --- a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_test.py +++ b/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(