Fix: Fix issue where Universe Domain is not correctly set. (#2423)

PhongChuong · gcf-owl-bot[bot] · web-flow · commit b7ebd73754fe · 2024-02-29T14:27:38.000-05:00
* Fix: Fix issue where Universe Domain is not correctly set. This PR fixes an issue where the universe domain is not correctly set. This also adds integration tests to verify the change. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Add universe domain integration tests for v1beta1 * Fix issue with universe domain for v1beta and add functional integration test for universe domain. * Fix integration test failure due to commented out code. Also close newly created clients. --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java
@@ -70,6 +70,7 @@ public static EnhancedBigQueryReadStub create(
     // Configure the base settings.
     BigQueryReadStubSettings.Builder baseSettingsBuilder =
         BigQueryReadStubSettings.newBuilder()
+            .setUniverseDomain(settings.getUniverseDomain())
             .setTransportChannelProvider(settings.getTransportChannelProvider())
             .setEndpoint(settings.getEndpoint())
             .setHeaderProvider(settings.getHeaderProvider())
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStubSettings.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStubSettings.java
@@ -159,7 +159,6 @@ protected Builder(ClientContext clientContext) {
 
       // Defaults provider
       BigQueryReadStubSettings.Builder baseDefaults = BigQueryReadStubSettings.newBuilder();
-      setEndpoint(baseDefaults.getEndpoint());
       setTransportChannelProvider(defaultTransportChannelProvider());
       setCredentialsProvider(baseDefaults.getCredentialsProvider());
       setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStub.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStub.java
@@ -74,8 +74,8 @@ public static EnhancedBigQueryStorageStub create(
     // Configure the base settings.
     BigQueryStorageStubSettings.Builder baseSettingsBuilder =
         BigQueryStorageStubSettings.newBuilder()
+            .setUniverseDomain(settings.getUniverseDomain())
             .setTransportChannelProvider(settings.getTransportChannelProvider())
-            .setEndpoint(settings.getEndpoint())
             .setHeaderProvider(settings.getHeaderProvider())
             .setCredentialsProvider(settings.getCredentialsProvider())
             .setStreamWatchdogCheckInterval(settings.getStreamWatchdogCheckInterval())
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStubSettings.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStubSettings.java
@@ -186,7 +186,6 @@ protected Builder(ClientContext clientContext) {
 
       // Defaults provider
       BigQueryStorageStubSettings.Builder baseDefaults = BigQueryStorageStubSettings.newBuilder();
-      setEndpoint(baseDefaults.getEndpoint());
       setTransportChannelProvider(defaultTransportChannelProvider());
       setCredentialsProvider(baseDefaults.getCredentialsProvider());
       setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/EnhancedBigQueryReadStub.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/EnhancedBigQueryReadStub.java
@@ -70,6 +70,7 @@ public static EnhancedBigQueryReadStub create(
     // Configure the base settings.
     BigQueryReadStubSettings.Builder baseSettingsBuilder =
         BigQueryReadStubSettings.newBuilder()
+            .setUniverseDomain(settings.getUniverseDomain())
             .setTransportChannelProvider(settings.getTransportChannelProvider())
             .setEndpoint(settings.getEndpoint())
             .setHeaderProvider(settings.getHeaderProvider())
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/EnhancedBigQueryReadStubSettings.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/EnhancedBigQueryReadStubSettings.java
@@ -159,7 +159,6 @@ protected Builder(ClientContext clientContext) {
 
       // Defaults provider
       BigQueryReadStubSettings.Builder baseDefaults = BigQueryReadStubSettings.newBuilder();
-      setEndpoint(baseDefaults.getEndpoint());
       setTransportChannelProvider(defaultTransportChannelProvider());
       setCredentialsProvider(baseDefaults.getCredentialsProvider());
       setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageTest.java
@@ -16,15 +16,20 @@
 
 package com.google.cloud.bigquery.storage.v1.it;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
+import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.core.InstantiatingExecutorProvider;
 import com.google.api.gax.rpc.ServerStream;
+import com.google.api.gax.rpc.UnauthenticatedException;
+import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.RetryOption;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.bigquery.BigQuery;
@@ -54,7 +59,9 @@
 import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.Timestamp;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.math.BigDecimal;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -82,7 +89,6 @@
 
 /** Integration tests for BigQuery Storage API. */
 public class ITBigQueryStorageTest {
-
   private static final Logger LOG = Logger.getLogger(ITBigQueryStorageTest.class.getName());
   private static final String DATASET = RemoteBigQueryHelper.generateDatasetName();
   private static final String DESCRIPTION = "BigQuery Storage Java client test dataset";
@@ -91,6 +97,66 @@ public class ITBigQueryStorageTest {
   private static String parentProjectId;
   private static BigQuery bigquery;
 
+  private static final String FAKE_JSON_CRED_WITH_GOOGLE_DOMAIN =
+      "{\n"
+          + "  \"private_key_id\": \"somekeyid\",\n"
+          + "  \"private_key\": \"-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggS"
+          + "kAgEAAoIBAQC+K2hSuFpAdrJI\\nnCgcDz2M7t7bjdlsadsasad+fvRSW6TjNQZ3p5LLQY1kSZRqBqylRkzteMOyHg"
+          + "aR\\n0Pmxh3ILCND5men43j3h4eDbrhQBuxfEMalkG92sL+PNQSETY2tnvXryOvmBRwa/\\nQP/9dJfIkIDJ9Fw9N4"
+          + "Bhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nknddadwkwewcVxHFhcZJO+XWf6ofLUXpRwiTZakGMn8EE1uVa2"
+          + "LgczOjwWHGi99MFjxSer5m9\\n1tCa3/KEGKiS/YL71JvjwX3mb+cewlkcmweBKZHM2JPTk0ZednFSpVZMtycjkbLa"
+          + "\\ndYOS8V85AgMBewECggEBAKksaldajfDZDV6nGqbFjMiizAKJolr/M3OQw16K6o3/\\n0S31xIe3sSlgW0+UbYlF"
+          + "4U8KifhManD1apVSC3csafaspP4RZUHFhtBywLO9pR5c\\nr6S5aLp+gPWFyIp1pfXbWGvc5VY/v9x7ya1VEa6rXvL"
+          + "sKupSeWAW4tMj3eo/64ge\\nsdaceaLYw52KeBYiT6+vpsnYrEkAHO1fF/LavbLLOFJmFTMxmsNaG0tuiJHgjshB\\"
+          + "n82DpMCbXG9YcCgI/DbzuIjsdj2JC1cascSP//3PmefWysucBQe7Jryb6NQtASmnv\\nCdDw/0jmZTEjpe4S1lxfHp"
+          + "lAhHFtdgYTvyYtaLZiVVkCgYEA8eVpof2rceecw/I6\\n5ng1q3Hl2usdWV/4mZMvR0fOemacLLfocX6IYxT1zA1FF"
+          + "JlbXSRsJMf/Qq39mOR2\\nSpW+hr4jCoHeRVYLgsbggtrevGmILAlNoqCMpGZ6vDmJpq6ECV9olliDvpPgWOP+\\nm"
+          + "YPDreFBGxWvQrADNbRt2dmGsrsCgYEAyUHqB2wvJHFqdmeBsaacewzV8x9WgmeX\\ngUIi9REwXlGDW0Mz50dxpxcK"
+          + "CAYn65+7TCnY5O/jmL0VRxU1J2mSWyWTo1C+17L0\\n3fUqjxL1pkefwecxwecvC+gFFYdJ4CQ/MHHXU81Lwl1iWdF"
+          + "Cd2UoGddYaOF+KNeM\\nHC7cmqra+JsCgYEAlUNywzq8nUg7282E+uICfCB0LfwejuymR93CtsFgb7cRd6ak\\nECR"
+          + "8FGfCpH8ruWJINllbQfcHVCX47ndLZwqv3oVFKh6pAS/vVI4dpOepP8++7y1u\\ncoOvtreXCX6XqfrWDtKIvv0vjl"
+          + "HBhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nkndj5uNl5SiuVxHFhcZJO+XWf6ofLUregtevZakGMn8EE1uVa"
+          + "2AY7eafmoU/nZPT\\n00YB0TBATdCbn/nBSuKDESkhSg9s2GEKQZG5hBmL5uCMfo09z3SfxZIhJdlerreP\\nJ7gSi"
+          + "dI12N+EZxYd4xIJh/HFDgp7RRO87f+WJkofMQKBgGTnClK1VMaCRbJZPriw\\nEfeFCoOX75MxKwXs6xgrw4W//AYG"
+          + "GUjDt83lD6AZP6tws7gJ2IwY/qP7+lyhjEqN\\nHtfPZRGFkGZsdaksdlaksd323423d+15/UvrlRSFPNj1tWQmNKk"
+          + "XyRDW4IG1Oa2p\\nrALStNBx5Y9t0/LQnFI4w3aG\\n-----END PRIVATE KEY-----\\n\",\n"
+          + "  \"project_id\": \"someprojectid\",\n"
+          + "  \"client_email\": \"someclientid@developer.gserviceaccount.com\",\n"
+          + "  \"client_id\": \"someclientid.apps.googleusercontent.com\",\n"
+          + "  \"type\": \"service_account\",\n"
+          + "  \"universe_domain\": \"googleapis.com\"\n"
+          + "}";
+
+  private static final String FAKE_JSON_CRED_WITH_INVALID_DOMAIN =
+      "{\n"
+          + "  \"private_key_id\": \"somekeyid\",\n"
+          + "  \"private_key\": \"-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggS"
+          + "kAgEAAoIBAQC+K2hSuFpAdrJI\\nnCgcDz2M7t7bjdlsadsasad+fvRSW6TjNQZ3p5LLQY1kSZRqBqylRkzteMOyHg"
+          + "aR\\n0Pmxh3ILCND5men43j3h4eDbrhQBuxfEMalkG92sL+PNQSETY2tnvXryOvmBRwa/\\nQP/9dJfIkIDJ9Fw9N4"
+          + "Bhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nknddadwkwewcVxHFhcZJO+XWf6ofLUXpRwiTZakGMn8EE1uVa2"
+          + "LgczOjwWHGi99MFjxSer5m9\\n1tCa3/KEGKiS/YL71JvjwX3mb+cewlkcmweBKZHM2JPTk0ZednFSpVZMtycjkbLa"
+          + "\\ndYOS8V85AgMBewECggEBAKksaldajfDZDV6nGqbFjMiizAKJolr/M3OQw16K6o3/\\n0S31xIe3sSlgW0+UbYlF"
+          + "4U8KifhManD1apVSC3csafaspP4RZUHFhtBywLO9pR5c\\nr6S5aLp+gPWFyIp1pfXbWGvc5VY/v9x7ya1VEa6rXvL"
+          + "sKupSeWAW4tMj3eo/64ge\\nsdaceaLYw52KeBYiT6+vpsnYrEkAHO1fF/LavbLLOFJmFTMxmsNaG0tuiJHgjshB\\"
+          + "n82DpMCbXG9YcCgI/DbzuIjsdj2JC1cascSP//3PmefWysucBQe7Jryb6NQtASmnv\\nCdDw/0jmZTEjpe4S1lxfHp"
+          + "lAhHFtdgYTvyYtaLZiVVkCgYEA8eVpof2rceecw/I6\\n5ng1q3Hl2usdWV/4mZMvR0fOemacLLfocX6IYxT1zA1FF"
+          + "JlbXSRsJMf/Qq39mOR2\\nSpW+hr4jCoHeRVYLgsbggtrevGmILAlNoqCMpGZ6vDmJpq6ECV9olliDvpPgWOP+\\nm"
+          + "YPDreFBGxWvQrADNbRt2dmGsrsCgYEAyUHqB2wvJHFqdmeBsaacewzV8x9WgmeX\\ngUIi9REwXlGDW0Mz50dxpxcK"
+          + "CAYn65+7TCnY5O/jmL0VRxU1J2mSWyWTo1C+17L0\\n3fUqjxL1pkefwecxwecvC+gFFYdJ4CQ/MHHXU81Lwl1iWdF"
+          + "Cd2UoGddYaOF+KNeM\\nHC7cmqra+JsCgYEAlUNywzq8nUg7282E+uICfCB0LfwejuymR93CtsFgb7cRd6ak\\nECR"
+          + "8FGfCpH8ruWJINllbQfcHVCX47ndLZwqv3oVFKh6pAS/vVI4dpOepP8++7y1u\\ncoOvtreXCX6XqfrWDtKIvv0vjl"
+          + "HBhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nkndj5uNl5SiuVxHFhcZJO+XWf6ofLUregtevZakGMn8EE1uVa"
+          + "2AY7eafmoU/nZPT\\n00YB0TBATdCbn/nBSuKDESkhSg9s2GEKQZG5hBmL5uCMfo09z3SfxZIhJdlerreP\\nJ7gSi"
+          + "dI12N+EZxYd4xIJh/HFDgp7RRO87f+WJkofMQKBgGTnClK1VMaCRbJZPriw\\nEfeFCoOX75MxKwXs6xgrw4W//AYG"
+          + "GUjDt83lD6AZP6tws7gJ2IwY/qP7+lyhjEqN\\nHtfPZRGFkGZsdaksdlaksd323423d+15/UvrlRSFPNj1tWQmNKk"
+          + "XyRDW4IG1Oa2p\\nrALStNBx5Y9t0/LQnFI4w3aG\\n-----END PRIVATE KEY-----\\n\",\n"
+          + "  \"project_id\": \"someprojectid\",\n"
+          + "  \"client_email\": \"someclientid@developer.gserviceaccount.com\",\n"
+          + "  \"client_id\": \"someclientid.apps.googleusercontent.com\",\n"
+          + "  \"type\": \"service_account\",\n"
+          + "  \"universe_domain\": \"fake.domain\"\n"
+          + "}";
+
   @BeforeClass
   public static void beforeClass() throws IOException {
     client = BigQueryReadClient.create();
@@ -859,6 +925,147 @@ public void testSimpleReadWithBackgroundExecutorProvider() throws IOException {
     assertEquals(164_656, rowCount);
   }
 
+  @Test
+  public void testUniverseDomainWithInvalidUniverseDomain() throws IOException {
+    BigQueryReadSettings bigQueryReadSettings =
+        BigQueryReadSettings.newBuilder()
+            .setCredentialsProvider(
+                FixedCredentialsProvider.create(loadCredentials(FAKE_JSON_CRED_WITH_GOOGLE_DOMAIN)))
+            .setUniverseDomain("invalid.domain")
+            .build();
+    BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);
+
+    String table =
+        BigQueryResource.FormatTableResource(
+            /* projectId = */ "bigquery-public-data",
+            /* datasetId = */ "samples",
+            /* tableId = */ "shakespeare");
+
+    try {
+      localClient.createReadSession(
+          /* parent = */ parentProjectId,
+          /* readSession = */ ReadSession.newBuilder()
+              .setTable(table)
+              .setDataFormat(DataFormat.AVRO)
+              .build(),
+          /* maxStreamCount = */ 1);
+      fail("RPCs to invalid universe domain should fail");
+    } catch (UnauthenticatedException e) {
+      assertThat(
+              (e.getMessage()
+                  .contains("does not match the universe domain found in the credentials")))
+          .isTrue();
+    }
+    localClient.close();
+  }
+
+  @Test
+  public void testInvalidUniverseDomainWithMismatchCredentials() throws IOException {
+    BigQueryReadSettings bigQueryReadSettings =
+        BigQueryReadSettings.newBuilder()
+            .setCredentialsProvider(
+                FixedCredentialsProvider.create(
+                    loadCredentials(FAKE_JSON_CRED_WITH_INVALID_DOMAIN)))
+            .setUniverseDomain("invalid.domain")
+            .build();
+    BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);
+
+    String table =
+        BigQueryResource.FormatTableResource(
+            /* projectId = */ "bigquery-public-data",
+            /* datasetId = */ "samples",
+            /* tableId = */ "shakespeare");
+
+    try {
+      ReadSession session =
+          localClient.createReadSession(
+              /* parent = */ parentProjectId,
+              /* readSession = */ ReadSession.newBuilder()
+                  .setTable(table)
+                  .setDataFormat(DataFormat.AVRO)
+                  .build(),
+              /* maxStreamCount = */ 1);
+      fail("RPCs to invalid universe domain should fail");
+    } catch (UnauthenticatedException e) {
+      assertThat(
+              (e.getMessage()
+                  .contains("does not match the universe domain found in the credentials")))
+          .isTrue();
+    }
+    localClient.close();
+  }
+
+  @Test
+  public void testUniverseDomainWithMatchingDomain() throws IOException {
+    // Test a valid domain using the default credentials and Google default universe domain.
+    BigQueryReadSettings bigQueryReadSettings =
+        BigQueryReadSettings.newBuilder().setUniverseDomain("googleapis.com").build();
+    BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);
+
+    String table =
+        BigQueryResource.FormatTableResource(
+            /* projectId = */ "bigquery-public-data",
+            /* datasetId = */ "samples",
+            /* tableId = */ "shakespeare");
+
+    ReadSession session =
+        localClient.createReadSession(
+            /* parent = */ parentProjectId,
+            /* readSession = */ ReadSession.newBuilder()
+                .setTable(table)
+                .setDataFormat(DataFormat.AVRO)
+                .build(),
+            /* maxStreamCount = */ 1);
+
+    ReadRowsRequest readRowsRequest =
+        ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build();
+
+    long rowCount = 0;
+    ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
+    for (ReadRowsResponse response : stream) {
+      rowCount += response.getRowCount();
+    }
+
+    assertEquals(164_656, rowCount);
+    localClient.close();
+  }
+
+  public void testUniverseDomain() throws IOException {
+    // This test is not yet part presubmit integration test as it requires the apis-tpclp.goog
+    // universe domain credentials.
+    // Test a valid read session in the universe domain gdutst.
+    BigQueryReadSettings bigQueryReadSettings =
+        BigQueryReadSettings.newBuilder().setUniverseDomain("apis-tpclp.goog").build();
+    BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);
+
+    String table =
+        BigQueryResource.FormatTableResource(
+            /* projectId = */ "google-tpc-testing-environment:cloudsdk-test-project",
+            /* datasetId = */ "tpc_demo_dataset",
+            /* tableId = */ "new_table");
+
+    ReadSession session =
+        localClient.createReadSession(
+            /* parent = */ parentProjectId,
+            /* readSession = */ ReadSession.newBuilder()
+                .setTable(table)
+                .setDataFormat(DataFormat.AVRO)
+                .build(),
+            /* maxStreamCount = */ 1);
+
+    ReadRowsRequest readRowsRequest =
+        ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build();
+
+    long rowCount = 0;
+    ServerStream<ReadRowsResponse> stream = localClient.readRowsCallable().call(readRowsRequest);
+    for (ReadRowsResponse response : stream) {
+      rowCount += response.getRowCount();
+    }
+
+    assertEquals(1, rowCount);
+    localClient.close();
+  }
+
   /**
    * Reads to the specified row offset within the stream. If the stream does not have the desired
    * rows to read, it will read all of them.
@@ -1015,4 +1222,14 @@ private Job RunQueryJobAndExpectSuccess(QueryJobConfiguration configuration)
 
     return completedJob;
   }
+
+  static GoogleCredentials loadCredentials(String credentialFile) {
+    try {
+      InputStream keyStream = new ByteArrayInputStream(credentialFile.getBytes());
+      return GoogleCredentials.fromStream(keyStream);
+    } catch (IOException e) {
+      fail("Couldn't create fake JSON credentials.");
+    }
+    return null;
+  }
 }
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java
