[Python] Fix GCS client creation in non-default universes

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()

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()

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.

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([