Merge pull request #32362 from iht/fix_gcsio_options

Fix using `GcsIO` from a `DoFn`in Python
apache · Aug 30, 2024 · fd8d146 · fd8d146
2 parents 322c47d + aac398d
commit fd8d146
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -145,10 +145,13 @@ def __init__(self, storage_client=None, pipeline_options=None):
       pipeline_options = PipelineOptions.from_dictionary(pipeline_options)
     if storage_client is None:
       storage_client = create_storage_client(pipeline_options)
-    self.enable_read_bucket_metric = pipeline_options.get_all_options(
-    )['enable_bucket_read_metric_counter']
-    self.enable_write_bucket_metric = pipeline_options.get_all_options(
-    )['enable_bucket_write_metric_counter']
+
+    google_cloud_options = pipeline_options.view_as(GoogleCloudOptions)
+    self.enable_read_bucket_metric = getattr(
+        google_cloud_options, 'enable_bucket_read_metric_counter', False)
+    self.enable_write_bucket_metric = getattr(
+        google_cloud_options, 'enable_bucket_write_metric_counter', False)
+
     self.client = storage_client
     self._rewrite_cb = None
     self.bucket_to_project_number = {}