From 357619d1fd5679c9bc64d28c54f5e428f5014295 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Sun, 25 Oct 2026 21:41:09 -0400 Subject: [PATCH] [Java] Honor non-default universe domains in GcsUtilV2 java-storage ignores GOOGLE_CLOUD_UNIVERSE_DOMAIN and the credentials' universe domain, so GcsUtilV2 always targeted storage.googleapis.com. When no gcsEndpoint is set, resolve the universe domain (env var, then credentials) and set it on StorageOptions. Add unit tests. --- .../sdk/extensions/gcp/util/GcsUtilV2.java | 42 +++++++++++ .../extensions/gcp/util/GcsUtilV2Test.java | 75 +++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java index c0f40a9aec4b..39853aa57363 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java @@ -109,6 +109,9 @@ class GcsUtilV2 { private static final org.slf4j.Logger LOG = org.slf4j.LoggerFactory.getLogger(GcsUtilV2.class); + /** Environment variable that selects the Google Cloud universe domain. */ + @VisibleForTesting static final String UNIVERSE_DOMAIN_ENV_VAR = "GOOGLE_CLOUD_UNIVERSE_DOMAIN"; + public static class GcsUtilFactory implements DefaultValueFactory { @Override public GcsUtilV2 create(PipelineOptions options) { @@ -208,6 +211,16 @@ Storage storage() { String endpoint = gcsOptions.getGcsEndpoint(); if (endpoint != null) { storageOptionsBuilder.setHost(rootUrlOf(endpoint)); + } else { + // java-storage only takes the universe domain from StorageOptions; unlike the Google API + // client used by GcsUtilV1, it ignores GOOGLE_CLOUD_UNIVERSE_DOMAIN and the credentials' + // universe domain. Set it explicitly so that non-default universes (e.g. Trusted Partner + // Cloud) resolve to their own storage endpoint instead of storage.googleapis.com. + @Nullable String universeDomain = + resolveUniverseDomain(System.getenv(UNIVERSE_DOMAIN_ENV_VAR), credentials); + if (universeDomain != null) { + storageOptionsBuilder.setUniverseDomain(universeDomain); + } } // Use the process-wide shared HTTP transport rather than letting each client create its own, so @@ -407,6 +420,35 @@ private static String rootUrlOf(String endpoint) { } } + /** + * Returns the non-default universe domain to use for the storage client, or null for the default + * universe ({@code googleapis.com}). + * + *

{@code envUniverseDomain} (the value of {@code GOOGLE_CLOUD_UNIVERSE_DOMAIN}) takes + * precedence; otherwise the credentials' universe domain is used. This matches how the Google API + * client libraries, and Beam's Python SDK, resolve the universe domain. + */ + @VisibleForTesting + static @Nullable String resolveUniverseDomain( + @Nullable String envUniverseDomain, @Nullable Credentials credentials) { + String universeDomain = envUniverseDomain == null ? "" : envUniverseDomain.trim(); + if (universeDomain.isEmpty() && credentials != null) { + try { + universeDomain = credentials.getUniverseDomain(); + } catch (IOException e) { + LOG.warn( + "Unable to determine the universe domain of the credentials; using the default.", e); + return null; + } + } + if (universeDomain == null + || universeDomain.isEmpty() + || Credentials.GOOGLE_DEFAULT_UNIVERSE.equals(universeDomain)) { + return null; + } + return universeDomain; + } + // AccessDeniedException/FileAlreadyExistsException permit a null "other" argument, and these // exceptions permit a null detail message, but the JDK constructor stubs are not annotated for // nullness. diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java index 9335ab6f56e0..df0d2678eaac 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java @@ -25,6 +25,7 @@ import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.never; @@ -223,6 +224,80 @@ public void testInvalidGcsEndpointIsRejected() { assertThrows(IllegalArgumentException.class, () -> new GcsUtilV2(options)); } + /** A test credential in the given universe domain. */ + private static class UniverseCredential extends TestCredential { + private final String universeDomain; + + UniverseCredential(String universeDomain) { + this.universeDomain = universeDomain; + } + + @Override + public String getUniverseDomain() { + return universeDomain; + } + } + + /** A test credential whose universe domain cannot be determined. */ + private static class FailingUniverseCredential extends TestCredential { + @Override + public String getUniverseDomain() throws IOException { + throw new IOException("no universe domain"); + } + } + + @Test + public void testResolveUniverseDomainPrefersEnvVar() { + assertEquals( + "apis-tpczero.goog", + GcsUtilV2.resolveUniverseDomain( + " apis-tpczero.goog ", new UniverseCredential("other.goog"))); + } + + @Test + public void testResolveUniverseDomainFallsBackToCredentials() { + assertEquals( + "apis-tpczero.goog", + GcsUtilV2.resolveUniverseDomain(null, new UniverseCredential("apis-tpczero.goog"))); + assertEquals( + "apis-tpczero.goog", + GcsUtilV2.resolveUniverseDomain("", new UniverseCredential("apis-tpczero.goog"))); + } + + @Test + public void testResolveUniverseDomainReturnsNullForDefaultUniverse() { + assertNull(GcsUtilV2.resolveUniverseDomain(null, null)); + assertNull(GcsUtilV2.resolveUniverseDomain(Credentials.GOOGLE_DEFAULT_UNIVERSE, null)); + assertNull(GcsUtilV2.resolveUniverseDomain(null, new TestCredential())); + } + + @Test + public void testResolveUniverseDomainIgnoresCredentialErrors() { + assertNull(GcsUtilV2.resolveUniverseDomain(null, new FailingUniverseCredential())); + } + + @Test + public void testCredentialsUniverseDomainIsUsedWithoutGcsEndpoint() { + // Skip if the environment variable is set, since it takes precedence over the credentials. + assumeTrue(System.getenv(GcsUtilV2.UNIVERSE_DOMAIN_ENV_VAR) == null); + GcsOptions options = gcsOptions(); + options.setGcpCredential(new UniverseCredential("example.goog")); + + StorageOptions storageOptions = storageOptionsOf(new GcsUtilV2(options)); + + assertEquals("example.goog", storageOptions.getUniverseDomain()); + assertEquals("https://storage.example.goog/", storageOptions.getResolvedHost("storage")); + } + + @Test + public void testGcsEndpointTakesPrecedenceOverUniverseDomain() { + GcsOptions options = gcsOptions(); + options.setGcpCredential(new UniverseCredential("example.goog")); + options.setGcsEndpoint("https://storage.example.com/storage/v1/"); + + assertEquals("https://storage.example.com", storageOptionsOf(new GcsUtilV2(options)).getHost()); + } + // --------------------------------------------------------------------------------------------- // HTTP metrics binding // ---------------------------------------------------------------------------------------------