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 // ---------------------------------------------------------------------------------------------