This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new c31bbe9e209 [Python] Catch PermissionDenied when checking the feature
store in VertexAIFeatureStoreEnrichmentHandler (#40463)
c31bbe9e209 is described below
commit c31bbe9e2092367545f4f48f2e8c245c283e8962
Author: Divyansh Yadav <[email protected]>
AuthorDate: Wed Oct 14 23:56:51 2026 +0530
[Python] Catch PermissionDenied when checking the feature store in
VertexAIFeatureStoreEnrichmentHandler (#40463)
* [Python] Catch PermissionDenied when checking the feature store in
VertexAIFeatureStoreEnrichmentHandler
* Move CHANGES.md entry to avoid merge conflict with master
---
CHANGES.md | 1 +
.../enrichment_handlers/vertex_ai_feature_store.py | 19 +++++-----
.../vertex_ai_feature_store_test.py | 40 ++++++++++++++++++++++
3 files changed, 51 insertions(+), 9 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 7693e65d78a..02f3818b1a8 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -99,6 +99,7 @@
* (Go) Fixed a deadlock recreating a data channel while holding the channel
lock ([#40414](https://github.com/apache/beam/issues/40414)).
* (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 c364378ea8a..88d2f3cc6ae 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 logging
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 @@ class
VertexAIFeatureStoreEnrichmentHandler(EnrichmentSourceHandler[beam.Row,
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 211529be4dc..233f0e11da6 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 @@ class TestVertexAIFeatureStoreHandlerInit(unittest.TestCase):
client_options={'api_endpoint': '[email protected]'},
)
+ @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='[email protected]',
+ 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='[email protected]',
+ 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(
