This is an automated email from the ASF dual-hosted git repository.
shunping 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 d746d747ef3 [Python] Fix GCS client creation in non-default universes
(#40674)
d746d747ef3 is described below
commit d746d747ef3fe88c1591d74bc8e0659fdc107368
Author: Shunping Huang <[email protected]>
AuthorDate: Tue Oct 20 09:36:24 2026 -0400
[Python] Fix GCS client creation in non-default universes (#40674)
---
sdks/python/apache_beam/io/gcp/gcsio.py | 21 ++++++++-
sdks/python/apache_beam/io/gcp/gcsio_test.py | 46 +++++++++++++++++++
.../python/apache_beam/options/pipeline_options.py | 8 ++--
.../apache_beam/options/pipeline_options_test.py | 52 ++++++++++++++++++++++
4 files changed, 123 insertions(+), 4 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py
b/sdks/python/apache_beam/io/gcp/gcsio.py
index a673845811d..2bdc4ca7d7e 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -30,6 +30,7 @@ https://github.com/apache/beam/blob/master/sdks/python/OWNERS
# pytype: skip-file
import logging
+import os
import re
import time
from typing import Optional
@@ -69,6 +70,8 @@ _LOGGER = logging.getLogger(__name__)
DEFAULT_READ_BUFFER_SIZE = 16 * 1024 * 1024
DEFAULT_WRITE_BUFFER_SIZE = 16 * 1024 * 1024
+_DEFAULT_UNIVERSE_DOMAIN = 'googleapis.com'
+
# Writes are performed as resumable uploads, which require the chunk size to be
# a multiple of 256 KiB.
WRITE_BUFFER_SIZE_MULTIPLE = 256 * 1024
@@ -205,10 +208,26 @@ def create_storage_client(pipeline_options,
use_credentials=True):
if google_cloud_options.gcs_custom_audit_entries is not None:
extra_headers.update(google_cloud_options.gcs_custom_audit_entries)
+ google_auth_credentials = credentials.get_google_auth_credentials()
+
+ # storage.Client only reads the universe domain from client_options; it
+ # ignores GOOGLE_CLOUD_UNIVERSE_DOMAIN and the credentials' universe
+ # domain, and fails if the latter differs from its default
+ # (googleapis.com). Set it explicitly to support non-default universes
+ # such as Trusted Partner Cloud (TPC).
+ universe_domain = (
+ os.environ.get('GOOGLE_CLOUD_UNIVERSE_DOMAIN') or
+ getattr(google_auth_credentials, 'universe_domain', None))
+ client_options = None
+ if (isinstance(universe_domain, str) and
+ universe_domain != _DEFAULT_UNIVERSE_DOMAIN):
+ client_options = {'universe_domain': universe_domain}
+
return storage.Client(
- credentials=credentials.get_google_auth_credentials(),
+ credentials=google_auth_credentials,
project=google_cloud_options.project,
client_info=beam_client_info,
+ client_options=client_options,
extra_headers=extra_headers)
else:
return storage.Client.create_anonymous_client()
diff --git a/sdks/python/apache_beam/io/gcp/gcsio_test.py
b/sdks/python/apache_beam/io/gcp/gcsio_test.py
index ddb94f9c17f..30e330e3afb 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio_test.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio_test.py
@@ -959,6 +959,52 @@ class TestGCSIO(unittest.TestCase):
self.gcs.is_soft_delete_enabled("gs://beam_without_soft_delete/tmp"))
+_TPC_UNIVERSE_DOMAIN = "example-tpc.goog"
+
+
[email protected](NotFound is None, 'GCP dependencies are not installed')
+# Disable mTLS auto-detection, which is incompatible with non-default
+# universe domains and may be enabled on developer machines.
[email protected](os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "false"})
+class TestCreateStorageClientUniverseDomain(unittest.TestCase):
+ """Tests for Trusted Partner Cloud (TPC) universe domain support."""
+ def _create_client(self, mock_get_service_credentials, universe_domain):
+ from apache_beam.internal.gcp.auth import _ApitoolsCredentialsAdapter
+ mock_get_service_credentials.return_value = _ApitoolsCredentialsAdapter(
+ _make_credentials("test-project", universe_domain=universe_domain))
+ return gcsio.create_storage_client(
+ PipelineOptions(["--project=test-project"]))
+
+ @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials')
+ def test_default_universe_domain(self, mock_get_service_credentials):
+ os.environ.pop("GOOGLE_CLOUD_UNIVERSE_DOMAIN", None)
+ client = self._create_client(
+ mock_get_service_credentials, _DEFAULT_UNIVERSE_DOMAIN)
+ self.assertEqual(client.universe_domain, _DEFAULT_UNIVERSE_DOMAIN)
+ self.assertEqual(
+ client._connection.API_BASE_URL, "https://storage.googleapis.com";)
+
+ @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials')
+ def test_universe_domain_from_env_var(self, mock_get_service_credentials):
+ os.environ["GOOGLE_CLOUD_UNIVERSE_DOMAIN"] = _TPC_UNIVERSE_DOMAIN
+ client = self._create_client(
+ mock_get_service_credentials, _TPC_UNIVERSE_DOMAIN)
+ self.assertEqual(client.universe_domain, _TPC_UNIVERSE_DOMAIN)
+ self.assertEqual(
+ client._connection.API_BASE_URL,
+ "https://storage.%s"; % _TPC_UNIVERSE_DOMAIN)
+
+ @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials')
+ def test_universe_domain_from_credentials(self,
mock_get_service_credentials):
+ os.environ.pop("GOOGLE_CLOUD_UNIVERSE_DOMAIN", None)
+ client = self._create_client(
+ mock_get_service_credentials, _TPC_UNIVERSE_DOMAIN)
+ self.assertEqual(client.universe_domain, _TPC_UNIVERSE_DOMAIN)
+ self.assertEqual(
+ client._connection.API_BASE_URL,
+ "https://storage.%s"; % _TPC_UNIVERSE_DOMAIN)
+
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
unittest.main()
diff --git a/sdks/python/apache_beam/options/pipeline_options.py
b/sdks/python/apache_beam/options/pipeline_options.py
index f0406e2505b..dca07da5146 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -1275,7 +1275,7 @@ class GoogleCloudOptions(PipelineOptions):
'Unable to check soft delete policy because GCP dependencies are '
'not installed.')
return
- if gcsio.GcsIO().is_soft_delete_enabled(gcs_path):
+ if gcsio.GcsIO(pipeline_options=self).is_soft_delete_enabled(gcs_path):
logger.log_first_n(
logging.WARN,
"Bucket %s used as %s has soft-delete policy enabled."
@@ -1289,8 +1289,10 @@ class GoogleCloudOptions(PipelineOptions):
arg_name,
n=1,
key="message")
- except ImportError:
- _LOGGER.warning('Unable to check soft delete policy due to import
error.')
+ except Exception as e: # pylint: disable=broad-except
+ # This check only emits a warning, so it must not block pipeline
+ # submission.
+ _LOGGER.warning('Unable to check soft delete policy: %s', e)
# If either temp or staging location has an issue, we use the valid one for
# both locations. If both are bad we return an error.
diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py
b/sdks/python/apache_beam/options/pipeline_options_test.py
index 493e42d5301..2814995b067 100644
--- a/sdks/python/apache_beam/options/pipeline_options_test.py
+++ b/sdks/python/apache_beam/options/pipeline_options_test.py
@@ -930,6 +930,58 @@ class PipelineOptionsTest(unittest.TestCase):
errors = options._handle_temp_and_staging_locations(validator)
self.assertEqual(errors, expected)
+ @unittest.skipIf(not has_gcsio, 'GCP dependencies are not installed')
+ @mock.patch.dict(
+ os.environ,
+ {
+ "GOOGLE_CLOUD_UNIVERSE_DOMAIN": "example-tpc.goog",
+ # mTLS auto-detection is incompatible with non-default universes.
+ "GOOGLE_API_USE_CLIENT_CERTIFICATE": "false",
+ })
+ @mock.patch('apache_beam.io.gcp.gcsio.GcsIO.get_bucket')
+ @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials')
+ def test_soft_delete_check_with_tpc_universe_domain(
+ self, mock_get_service_credentials, mock_get_bucket):
+ import google.auth.credentials
+
+ from apache_beam.internal.gcp.auth import _ApitoolsCredentialsAdapter
+ mock_get_service_credentials.return_value = _ApitoolsCredentialsAdapter(
+ mock.Mock(
+ spec=google.auth.credentials.Credentials,
+ universe_domain="example-tpc.goog"))
+ soft_delete_policy = mock_get_bucket.return_value.soft_delete_policy
+ soft_delete_policy.retention_duration_seconds = 0
+
+ options = GoogleCloudOptions([
+ '--project=myproject',
+ '--impersonate_service_account=s...@myproject.iam.gserviceaccount.com',
+ '--temp_location=gs://beam/tmp'
+ ])
+ # Must not raise a universe domain mismatch ValueError.
+ options._warn_if_soft_delete_policy_enabled('temp_location')
+
+ # The GCS client must be built from the user's pipeline options so that
+ # project and impersonation settings are honored.
+ used_options = mock_get_service_credentials.call_args[0][0].view_as(
+ GoogleCloudOptions)
+ self.assertEqual(used_options.project, 'myproject')
+ self.assertEqual(
+ used_options.impersonate_service_account,
+ '[email protected]')
+
+ @unittest.skipIf(not has_gcsio, 'GCP dependencies are not installed')
+ @mock.patch(
+ 'apache_beam.io.gcp.gcsio.GcsIO.__init__',
+ side_effect=ValueError('universe domain mismatch'))
+ def test_soft_delete_check_does_not_raise(self, unused_mock_init):
+ options = GoogleCloudOptions(
+ ['--project=myproject', '--temp_location=gs://beam/tmp'])
+ # The soft delete check only emits a warning; failures must not block
+ # pipeline submission.
+ with self.assertLogs('apache_beam.options.pipeline_options',
+ level='WARNING'):
+ options._warn_if_soft_delete_policy_enabled('temp_location')
+
def test_validation_good_stg_good_temp(self):
runner = MockRunners.DataflowRunner()
options = GoogleCloudOptions([
