[Java] Honor non-default universe domains in GcsUtilV2

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<GcsUtilV2> {
     @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}).
+   *
+   * <p>{@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.

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