test(bigquery): support regional endpoints dynamically in integration tests

Author: jinseopkim0

.kokoro/presubmit/bigquery-graalvm-native-presubmit.cfg

@@ -22,6 +22,11 @@ env_vars: {
   value: "gcloud-devel"
 }
 
+env_vars: {
+  key: "INTEGRATION_TEST_ARGS"
+  value: "-Dbigquery.endpoint=https://us-east7-bigquery.googleapis.com -Dbigquery.storage.endpoint=us-east7-bigquerystorage.googleapis.com:443"
+}
+
 env_vars: {
   key: "GOOGLE_APPLICATION_CREDENTIALS"
   value: "secret_manager/java-it-service-account"
@@ -40,5 +45,3 @@ env_vars: {
   key: "BUILD_SUBDIR"
   value: "java-bigquery"
 }
-
-

.kokoro/presubmit/bigquery-integration.cfg

@@ -22,6 +22,11 @@ env_vars: {
   value: "gcloud-devel"
 }
 
+env_vars: {
+  key: "INTEGRATION_TEST_ARGS"
+  value: "-Dbigquery.endpoint=https://us-east7-bigquery.googleapis.com -Dbigquery.storage.endpoint=us-east7-bigquerystorage.googleapis.com:443"
+}
+
 env_vars: {
   key: "GOOGLE_APPLICATION_CREDENTIALS"
   value: "secret_manager/java-it-service-account"

.kokoro/presubmit/bigquerystorage-graalvm-native-presubmit.cfg

@@ -44,5 +44,5 @@ env_vars: {
 
 env_vars: {
   key: "INTEGRATION_TEST_ARGS"
-  value: "-Dit.test=!ITBigQueryWrite*RetryTest -Dsurefire.failIfNoSpecifiedTests=false -Dfailsafe.failIfNoSpecifiedTests=false"
+  value: "-Dit.test=!ITBigQueryWrite*RetryTest -Dsurefire.failIfNoSpecifiedTests=false -Dfailsafe.failIfNoSpecifiedTests=false -Dbigquery.storage.endpoint=us-east7-bigquerystorage.googleapis.com:443 -Dbigquery.endpoint=https://us-east7-bigquery.googleapis.com"
 }

.kokoro/presubmit/bigquerystorage-integration.cfg

@@ -39,5 +39,5 @@ env_vars: {
 
 env_vars: {
   key: "INTEGRATION_TEST_ARGS"
-  value: "-Dit.test=!ITBigQueryWrite*RetryTest -Dsurefire.failIfNoSpecifiedTests=false -Dfailsafe.failIfNoSpecifiedTests=false"
+  value: "-Dit.test=!ITBigQueryWrite*RetryTest -Dsurefire.failIfNoSpecifiedTests=false -Dfailsafe.failIfNoSpecifiedTests=false -Dbigquery.storage.endpoint=us-east7-bigquerystorage.googleapis.com:443 -Dbigquery.endpoint=https://us-east7-bigquery.googleapis.com"
 }

java-bigquery/google-cloud-bigquery/pom.xml

@@ -15,6 +15,8 @@
   </parent>
   <properties>
     <site.installationModule>google-cloud-bigquery</site.installationModule>
+    <bigquery.endpoint></bigquery.endpoint>
+    <bigquery.storage.endpoint></bigquery.storage.endpoint>
   </properties>
   <dependencies>
     <dependency>
@@ -326,5 +328,22 @@
       </build>
 
     </profile>
+    <profile>
+      <id>native</id>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.graalvm.buildtools</groupId>
+            <artifactId>native-maven-plugin</artifactId>
+            <configuration>
+              <environmentVariables>
+                <BIGQUERY_ENDPOINT>${bigquery.endpoint}</BIGQUERY_ENDPOINT>
+                <BIGQUERY_STORAGE_ENDPOINT>${bigquery.storage.endpoint}</BIGQUERY_STORAGE_ENDPOINT>
+              </environmentVariables>
+            </configuration>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
   </profiles>
 </project>

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/testing/RemoteBigQueryHelper.java

@@ -105,7 +105,7 @@ public static RemoteBigQueryHelper create(String projectId, InputStream keyStrea
               .setProjectId(projectId)
               .setRetrySettings(retrySettings())
               .setTransportOptions(transportOptions);
-      String endpoint = System.getenv("BIGQUERY_ENDPOINT");
+      String endpoint = System.getProperty("bigquery.endpoint", System.getenv("BIGQUERY_ENDPOINT"));
       if (endpoint != null) {
         builder.setHost(endpoint);
       }
@@ -143,7 +143,7 @@ public static RemoteBigQueryHelper create(BigQueryOptions.Builder bigqueryOption
         bigqueryOptionsBuilder
             .setRetrySettings(retrySettings())
             .setTransportOptions(transportOptions);
-    String endpoint = System.getenv("BIGQUERY_ENDPOINT");
+    String endpoint = System.getProperty("bigquery.endpoint", System.getenv("BIGQUERY_ENDPOINT"));
     if (endpoint != null) {
       builder.setHost(endpoint);
     }
@@ -184,4 +184,16 @@ public static BigQueryHelperException translate(Exception ex) {
       return new BigQueryHelperException(ex.getMessage(), ex);
     }
   }
+
+  /**
+   * Helper to check if the provided BigQuery client is configured to target a regional endpoint.
+   */
+  public static boolean isRegionalEndpoint(BigQuery bigquery) {
+    if (bigquery == null || bigquery.getOptions() == null) {
+      return false;
+    }
+    String host = bigquery.getOptions().getHost();
+    return host != null
+        && (host.contains("-bigquery.googleapis.com") || host.contains(".rep.googleapis.com"));
+  }
 }

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -211,6 +211,7 @@
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Assumptions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
@@ -1271,6 +1272,8 @@ void testLosslessMaxTimestampIntegration() throws InterruptedException {
 
   @Test
   void testListDatasets() {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     Page<Dataset> datasets = bigquery.listDatasets("bigquery-public-data");
     Iterator<Dataset> iterator = datasets.iterateAll().iterator();
     Set<String> datasetNames = new HashSet<>();
@@ -2131,6 +2134,8 @@ void testCreateTableWithDefaultValueExpression() {
 
   @Test
   void testCreateAndUpdateTableWithPolicyTags() throws IOException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     // Set up policy tags in the datacatalog service
     try (PolicyTagManagerClient policyTagManagerClient = PolicyTagManagerClient.create()) {
       CreateTaxonomyRequest createTaxonomyRequest =
@@ -2515,6 +2520,8 @@ void testCreateExternalTable() throws InterruptedException {
 
   @Test
   void testSetPermExternalTableSchema() {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     String tableName = "test_create_external_table_perm";
     TableId tableId = TableId.of(DATASET, tableName);
     ExternalTableDefinition externalTableDefinition =
@@ -2926,6 +2933,8 @@ void testDeleteNonExistingTable() {
 
   @Test
   void testDeleteJob() {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     String query = "SELECT 17 as foo";
     QueryJobConfiguration config = QueryJobConfiguration.of(query);
     String jobName = "jobId_" + UUID.randomUUID().toString();
@@ -3188,6 +3197,8 @@ void testListAllTableData() {
 
   @Test
   void testListPageWithStartIndex() {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     String tableName = "midyear_population_agespecific";
     TableId tableId = TableId.of(PUBLIC_PROJECT, PUBLIC_DATASET, tableName);
     Table table = bigquery.getTable(tableId);
@@ -3659,6 +3670,8 @@ void testQueryStatistics() throws InterruptedException {
 
   @Test
   void testExecuteSelectDefaultConnectionSettings() throws SQLException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     // Use the default connection settings
     Connection connection = bigquery.createConnection();
     String query = "SELECT corpus FROM `bigquery-public-data.samples.shakespeare` GROUP BY corpus;";
@@ -3669,6 +3682,8 @@ void testExecuteSelectDefaultConnectionSettings() throws SQLException {
 
   @Test
   void testExecuteSelectWithReadApi() throws SQLException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     final int rowLimit = 5000;
     final String QUERY =
         "SELECT * FROM bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2017 LIMIT %s";
@@ -3699,6 +3714,8 @@ void testExecuteSelectWithReadApi() throws SQLException {
 
   @Test
   void testExecuteSelectWithFastQueryReadApi() throws SQLException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     final int rowLimit = 5000;
     final String QUERY =
         "SELECT * FROM bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2017 LIMIT %s";
@@ -4785,6 +4802,8 @@ void testProjectIDFastSQLQueryWithJobId() {
 
   @Test
   void testLocationFastSQLQueryWithJobId() throws InterruptedException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     TableId tableIdFastQueryUk = TableId.of(UK_DATASET, "fastquery_testing_table");
     DatasetInfo infoUK =
         DatasetInfo.newBuilder(UK_DATASET)
@@ -4960,6 +4979,8 @@ void testFastDDLQuery() throws InterruptedException {
 
   @Test
   void testFastQuerySlowDDL() throws InterruptedException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     String tableName = generateRandomName("test_table_fast_query_ddl_slow_");
     // This query take more than 10s to run and should fall back on the old query path
     String slowDdlQuery =
@@ -5061,6 +5082,8 @@ void testQuerySessionSupport() throws InterruptedException {
 
   @Test
   void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     TableId sessionTableId = TableId.of("_SESSION", "test_temp_destination_table_from_file");
 
     WriteChannelConfiguration configuration =
@@ -5288,6 +5311,8 @@ void testTransactionInfo() throws InterruptedException {
   /* TODO(prasmish): replicate the entire test case for executeSelect */
   @Test
   void testScriptStatistics() throws InterruptedException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     String script =
         "-- Declare a variable to hold names as an array.\n"
             + "DECLARE top_names ARRAY<STRING>;\n"
@@ -6541,6 +6566,8 @@ void testCancelJob() throws InterruptedException, TimeoutException {
 
   @Test
   void testCancelNonExistingJob() {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     assertFalse(bigquery.cancel("test_cancel_non_existing_job"));
   }
 
@@ -6677,6 +6704,8 @@ void testInsertWithDecimalTargetTypes()
 
   @Test
   void testLocation() throws Exception {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     String location = "EU";
     String wrongLocation = "US";
 
@@ -7443,6 +7472,8 @@ void testTableResultJobIdAndQueryId() throws InterruptedException {
 
   @Test
   void testStatelessQueriesWithLocation() throws Exception {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     // This test validates BigQueryOption location is used for stateless query by verifying that the
     // stateless query fails when the BigQueryOption location does not match the dataset location.
     String location = "EU";
@@ -7608,6 +7639,8 @@ void testInvalidUniverseDomainWithMismatchCredentials() {
 
   @Test
   void testUniverseDomainWithMatchingDomain() {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     // Test a valid domain using the default credentials and Google default universe domain.
     RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
     BigQueryOptions bigQueryOptions =
@@ -7701,6 +7734,8 @@ void testExternalMetadataCacheModeFailForNonBiglake() {
 
   @Test
   void testObjectTable() throws InterruptedException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     String tableName = generateRandomName("test_object_table");
     TableId tableId = TableId.of(DATASET, tableName);
 

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITNightlyBigQueryTest.java

@@ -69,6 +69,7 @@
 import java.util.logging.Logger;
 import org.apache.arrow.vector.util.JsonStringArrayList;
 import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Assumptions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
@@ -507,6 +508,8 @@ void testPositionalParams()
   // table-not-found exception. Ref: b/241134681 . This exception has been seen while reading data
   // in bulk
   void testForTableNotFound() throws SQLException {
+    // This test queries public datasets, which are not supported by regional endpoints.
+    Assumptions.assumeTrue(!RemoteBigQueryHelper.isRegionalEndpoint(bigquery));
     int recordCnt = 50000000; // 5Mil
     String query =
         String.format(

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java

@@ -94,12 +94,12 @@ public void testListDatasetsTraced() {
         assertEquals("GET", attrs.get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD));
         assertEquals("DatasetService", attrs.get(AttributeKey.stringKey("bq.rpc.service")));
         assertEquals("ListDatasets", attrs.get(AttributeKey.stringKey("bq.rpc.method")));
-        assertEquals(
-            "bigquery.googleapis.com", attrs.get(HttpTracingRequestInitializer.SERVER_ADDRESS));
+        assertEquals(getExpectedHost(), attrs.get(HttpTracingRequestInitializer.SERVER_ADDRESS));
         assertEquals(200L, attrs.get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE));
-        assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN));
+        assertEquals(getExpectedHost(), attrs.get(BigQueryTelemetryTracer.URL_DOMAIN));
         assertEquals(
-            "https://bigquery.googleapis.com/bigquery/v2/projects/"
+            getExpectedFullHost()
+                + "/bigquery/v2/projects/"
                 + bigqueryHelper.getOptions().getProjectId()
                 + "/datasets?prettyPrint=false",
             attrs.get(HttpTracingRequestInitializer.URL_FULL));
@@ -150,13 +150,13 @@ public void testGetDatasetNotFoundTraced() {
             "projects/{+projectId}/datasets/{+datasetId}",
             attrs.get(BigQueryTelemetryTracer.URL_TEMPLATE));
         assertEquals(
-            "https://bigquery.googleapis.com/bigquery/v2/projects/"
+            getExpectedFullHost()
+                + "/bigquery/v2/projects/"
                 + bigqueryHelper.getOptions().getProjectId()
                 + "/datasets/non_existent_dataset?prettyPrint=false",
             attrs.get(HttpTracingRequestInitializer.URL_FULL));
-        assertEquals(
-            "bigquery.googleapis.com", attrs.get(HttpTracingRequestInitializer.SERVER_ADDRESS));
-        assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN));
+        assertEquals(getExpectedHost(), attrs.get(HttpTracingRequestInitializer.SERVER_ADDRESS));
+        assertEquals(getExpectedHost(), attrs.get(BigQueryTelemetryTracer.URL_DOMAIN));
         assertEquals(
             "//bigquery.googleapis.com/projects/"
                 + bigqueryHelper.getOptions().getProjectId()
@@ -320,4 +320,35 @@ private void checkGeneralAttributes(Map<AttributeKey<?>, Object> attrs) {
         attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_ARTIFACT));
     assertNotNull(attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_VERSION));
   }
+
+  /**
+   * Returns the expected host header/domain to be matched in telemetry trace assertions.
+   * Dynamically strips protocol prefixes ("https://", "http://") from the configured BigQuery host
+   * options, falling back to the default "bigquery.googleapis.com" if not configured.
+   */
+  private static String getExpectedHost() {
+    String host = bigqueryHelper.getOptions().getHost();
+    if (host == null || host.isEmpty() || host.equals("https://www.googleapis.com")) {
+      return "bigquery.googleapis.com";
+    }
+    if (host.startsWith("https://")) {
+      return host.substring("https://".length());
+    }
+    if (host.startsWith("http://")) {
+      return host.substring("http://".length());
+    }
+    return host;
+  }
+
+  /**
+   * Returns the expected full URL host prefix (including the protocol scheme) for URL assertions.
+   * Defaults to "https://bigquery.googleapis.com" if the host option is not configured.
+   */
+  private static String getExpectedFullHost() {
+    String host = bigqueryHelper.getOptions().getHost();
+    if (host == null || host.isEmpty() || host.equals("https://www.googleapis.com")) {
+      return "https://bigquery.googleapis.com";
+    }
+    return host;
+  }
 }

java-bigquerystorage/google-cloud-bigquerystorage/pom.xml

@@ -15,6 +15,8 @@
   </parent>
   <properties>
     <site.installationModule>google-cloud-bigquerystorage</site.installationModule>
+    <bigquery.endpoint></bigquery.endpoint>
+    <bigquery.storage.endpoint></bigquery.storage.endpoint>
   </properties>
   <dependencyManagement>
     <dependencies>
@@ -409,6 +411,10 @@
                 <buildArg>--no-fallback</buildArg>
                 <buildArg>--no-server</buildArg>
               </buildArgs>
+              <environmentVariables>
+                <BIGQUERY_ENDPOINT>${bigquery.endpoint}</BIGQUERY_ENDPOINT>
+                <BIGQUERY_STORAGE_ENDPOINT>${bigquery.storage.endpoint}</BIGQUERY_STORAGE_ENDPOINT>
+              </environmentVariables>
             </configuration>
           </plugin>
         </plugins>