diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py index a673845811da..2bdc4ca7d7e4 100644 --- a/sdks/python/apache_beam/io/gcp/gcsio.py +++ b/sdks/python/apache_beam/io/gcp/gcsio.py @@ -30,6 +30,7 @@ # pytype: skip-file import logging +import os import re import time from typing import Optional @@ -69,6 +70,8 @@ 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 ddb94f9c17fa..30e330e3afba 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 @@ def test_is_soft_delete_enabled(self, mock_get_bucket): self.gcs.is_soft_delete_enabled("gs://beam_without_soft_delete/tmp")) +_TPC_UNIVERSE_DOMAIN = "example-tpc.goog" + + +@unittest.skipIf(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. +@mock.patch.dict(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 f0406e2505bc..dca07da5146e 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1275,7 +1275,7 @@ def _warn_if_soft_delete_policy_enabled(self, arg_name): '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 @@ def _warn_if_soft_delete_policy_enabled(self, arg_name): 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 493e42d53010..2814995b067b 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 @@ def _check_errors(self, options, validator, expected): 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=sa@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, + 'sa@myproject.iam.gserviceaccount.com') + + @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([