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..a30827c8d6f3 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 @@ -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.""" 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..fcdcd1a5bac4 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,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(