From 81baaa3554b1512b475f6f159e51d9990f8e5e1b Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 31 Jul 2026 19:44:11 -0400 Subject: [PATCH] test(bigquery): support regional endpoints dynamically in integration tests --- .../bigquery-graalvm-native-presubmit.cfg | 7 ++- .kokoro/presubmit/bigquery-integration.cfg | 5 ++ ...gquerystorage-graalvm-native-presubmit.cfg | 2 +- .../presubmit/bigquerystorage-integration.cfg | 2 +- java-bigquery/google-cloud-bigquery/pom.xml | 19 +++++++ .../testing/RemoteBigQueryHelper.java | 16 +++++- .../cloud/bigquery/it/ITBigQueryTest.java | 35 ++++++++++++ .../bigquery/it/ITNightlyBigQueryTest.java | 3 + .../bigquery/it/ITOpenTelemetryTest.java | 47 ++++++++++++--- .../google-cloud-bigquerystorage/pom.xml | 6 ++ .../it/ITBigQueryStorageReadClientTest.java | 40 ++++++++++--- .../it/ITBigQueryStorageWriteClientTest.java | 26 +++++++-- .../bigquery/storage/v1/it/util/Helper.java | 57 ++++++++++++++++++- 13 files changed, 236 insertions(+), 29 deletions(-) diff --git a/.kokoro/presubmit/bigquery-graalvm-native-presubmit.cfg b/.kokoro/presubmit/bigquery-graalvm-native-presubmit.cfg index 6fe61b51410a..ff79960ec1c2 100644 --- a/.kokoro/presubmit/bigquery-graalvm-native-presubmit.cfg +++ b/.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" } - - diff --git a/.kokoro/presubmit/bigquery-integration.cfg b/.kokoro/presubmit/bigquery-integration.cfg index 19770b533b3b..418a420a32ee 100644 --- a/.kokoro/presubmit/bigquery-integration.cfg +++ b/.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" diff --git a/.kokoro/presubmit/bigquerystorage-graalvm-native-presubmit.cfg b/.kokoro/presubmit/bigquerystorage-graalvm-native-presubmit.cfg index 597b40cafcf2..c772de8f6345 100644 --- a/.kokoro/presubmit/bigquerystorage-graalvm-native-presubmit.cfg +++ b/.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" } diff --git a/.kokoro/presubmit/bigquerystorage-integration.cfg b/.kokoro/presubmit/bigquerystorage-integration.cfg index 97bbe528e2b8..99aabeaa6103 100644 --- a/.kokoro/presubmit/bigquerystorage-integration.cfg +++ b/.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" } diff --git a/java-bigquery/google-cloud-bigquery/pom.xml b/java-bigquery/google-cloud-bigquery/pom.xml index fba380924aad..5db5f3331b53 100644 --- a/java-bigquery/google-cloud-bigquery/pom.xml +++ b/java-bigquery/google-cloud-bigquery/pom.xml @@ -15,6 +15,8 @@ google-cloud-bigquery + + @@ -326,5 +328,22 @@ + + native + + + + org.graalvm.buildtools + native-maven-plugin + + + ${bigquery.endpoint} + ${bigquery.storage.endpoint} + + + + + + diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/testing/RemoteBigQueryHelper.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/testing/RemoteBigQueryHelper.java index 49a3e5e5a482..f99dd4b01c5a 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/testing/RemoteBigQueryHelper.java +++ b/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")); + } } diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 3c9613d20758..b6d722f8fb67 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/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 datasets = bigquery.listDatasets("bigquery-public-data"); Iterator iterator = datasets.iterateAll().iterator(); Set 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;\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); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITNightlyBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITNightlyBigQueryTest.java index 641868203877..8392feabd299 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITNightlyBigQueryTest.java +++ b/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( diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java index 2b477754b531..87a3d2f67a7e 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java +++ b/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, 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; + } } diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml index 7370fb989ffd..9ffea7df847f 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml +++ b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml @@ -15,6 +15,8 @@ google-cloud-bigquerystorage + + @@ -409,6 +411,10 @@ --no-fallback --no-server + + ${bigquery.endpoint} + ${bigquery.storage.endpoint} + diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java index 6accfa6bd878..9b6b995ee5ee 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java @@ -117,6 +117,7 @@ import org.json.JSONArray; import org.json.JSONObject; import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assumptions; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -504,7 +505,7 @@ public CompletableResultCode shutdown() { @BeforeAll static void beforeAll() throws IOException, DescriptorValidationException, InterruptedException { - readClient = com.google.cloud.bigquery.storage.v1.it.util.Helper.createBigQueryReadClient(); + readClient = Helper.createBigQueryReadClient(); projectName = ServiceOptions.getDefaultProjectId(); parentProjectId = String.format("projects/%s", projectName); @@ -515,10 +516,19 @@ static void beforeAll() throws IOException, DescriptorValidationException, Inter RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); bigquery = bigqueryHelper.getOptions().getService(); - DatasetInfo datasetInfo = - DatasetInfo.newBuilder(/* datasetId bigquery= */ DATASET) - .setDescription(DESCRIPTION) - .build(); + DatasetInfo datasetInfo; + if (Helper.isBigQueryRegionalEndpoint()) { + datasetInfo = + DatasetInfo.newBuilder(/* datasetId bigquery= */ DATASET) + .setDescription(DESCRIPTION) + .setLocation(Helper.getBigQueryRegion()) + .build(); + } else { + datasetInfo = + DatasetInfo.newBuilder(/* datasetId bigquery= */ DATASET) + .setDescription(DESCRIPTION) + .build(); + } bigquery.create(datasetInfo); LOG.info("Created test dataset: " + DATASET); @@ -596,6 +606,8 @@ static void afterAll() throws InterruptedException { @Test void testSimpleReadAvro() { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); String table = BigQueryResource.formatTableResource( /* projectId= */ "bigquery-public-data", @@ -632,6 +644,8 @@ void testSimpleReadAvro() { @Test void testSimpleReadArrow() { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); String table = BigQueryResource.formatTableResource( /* projectId= */ "bigquery-public-data", @@ -948,6 +962,8 @@ void timestamp_readAvro() throws IOException { @Test void testSimpleReadAndResume() { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); String table = BigQueryResource.formatTableResource( /* projectId= */ "bigquery-public-data", @@ -991,6 +1007,8 @@ void testSimpleReadAndResume() { @Test void testFilter() throws IOException { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); String table = BigQueryResource.formatTableResource( /* projectId= */ "bigquery-public-data", @@ -1052,6 +1070,8 @@ record -> { @Test void testColumnSelection() throws IOException { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); String table = BigQueryResource.formatTableResource( /* projectId= */ "bigquery-public-data", @@ -1597,8 +1617,10 @@ void testStructAndArraySqlTypes() throws InterruptedException, IOException { @Test void testSimpleReadWithBackgroundExecutorProvider() throws IOException { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); BigQueryReadSettings bigQueryReadSettings = - com.google.cloud.bigquery.storage.v1.it.util.Helper.createBigQueryReadSettingsBuilder() + Helper.createBigQueryReadSettingsBuilder() .setBackgroundExecutorProvider( InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(14).build()) .build(); @@ -1718,6 +1740,8 @@ void testInvalidUniverseDomainWithMismatchCredentials() throws IOException { @Test void testUniverseDomainWithMatchingDomain() throws IOException { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); // Test a valid domain using the default credentials and Google default universe domain. BigQueryReadSettings bigQueryReadSettings = BigQueryReadSettings.newBuilder().setUniverseDomain("googleapis.com").build(); @@ -1753,6 +1777,8 @@ void testUniverseDomainWithMatchingDomain() throws IOException { @Test void testSimpleReadWithOtelTracing() throws IOException { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional endpoint does not support reading public datasets"); SdkTracerProvider tracerProvider = SdkTracerProvider.builder() .addSpanProcessor(SimpleSpanProcessor.create(new TestSpanExporter())) @@ -1761,7 +1787,7 @@ void testSimpleReadWithOtelTracing() throws IOException { OpenTelemetry otel = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); BigQueryReadSettings otelSettings = - com.google.cloud.bigquery.storage.v1.it.util.Helper.createBigQueryReadSettingsBuilder() + Helper.createBigQueryReadSettingsBuilder() .setEnableOpenTelemetryTracing(true) .setOpenTelemetryTracerProvider(tracerProvider) .build(); diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java index 2e80cdb443b6..01aa2031a0d2 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java @@ -93,6 +93,7 @@ import org.json.JSONArray; import org.json.JSONObject; 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.parallel.Execution; @@ -187,10 +188,10 @@ public StringWithSecondsNanos(String fooParam, long secondsParam, int nanosParam @BeforeAll static void beforeAll() throws IOException { - readClient = com.google.cloud.bigquery.storage.v1.it.util.Helper.createBigQueryReadClient(); + readClient = Helper.createBigQueryReadClient(); BigQueryWriteSettings settings = - com.google.cloud.bigquery.storage.v1.it.util.Helper.createBigQueryWriteSettingsBuilder() + Helper.createBigQueryWriteSettingsBuilder() .setHeaderProvider(USER_AGENT_HEADER_PROVIDER) .build(); writeClient = BigQueryWriteClient.create(settings); @@ -198,8 +199,17 @@ static void beforeAll() throws IOException { RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); bigquery = bigqueryHelper.getOptions().getService(); - DatasetInfo datasetInfo = - DatasetInfo.newBuilder(/* datasetId= */ DATASET).setDescription(DESCRIPTION).build(); + DatasetInfo datasetInfo; + if (Helper.isBigQueryRegionalEndpoint()) { + datasetInfo = + DatasetInfo.newBuilder(/* datasetId= */ DATASET) + .setDescription(DESCRIPTION) + .setLocation(Helper.getBigQueryRegion()) + .build(); + } else { + datasetInfo = + DatasetInfo.newBuilder(/* datasetId= */ DATASET).setDescription(DESCRIPTION).build(); + } bigquery.create(datasetInfo); LOG.info("Created test dataset: " + DATASET); tableInfo = @@ -373,6 +383,8 @@ ProtoRows createProtoRowsMixed(StringWithSecondsNanos[] messages) { @Test void testBatchWriteWithCommittedStreamEU() throws IOException, InterruptedException, ExecutionException { + Assumptions.assumeTrue( + !Helper.isRegionalEndpoint(), "Regional write client cannot write to EU table"); WriteStream writeStream = writeClient.createWriteStream( CreateWriteStreamRequest.newBuilder() @@ -2253,8 +2265,10 @@ void testMultiplexingMixedLocation() assertEquals(0L, response1.get().getAppendResult().getOffset().getValue()); assertEquals(0L, response2.get().getAppendResult().getOffset().getValue()); assertEquals(0L, response3.get().getAppendResult().getOffset().getValue()); - assertEquals("us", streamWriter1.getLocation()); - assertEquals("us", streamWriter2.getLocation()); + String expectedLocation = + Helper.isBigQueryRegionalEndpoint() ? Helper.getBigQueryRegion() : "us"; + assertEquals(expectedLocation, streamWriter1.getLocation()); + assertEquals(expectedLocation, streamWriter2.getLocation()); assertEquals("eu", streamWriter3.getLocation()); streamWriter1.close(); streamWriter2.close(); diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java index 8e66820b6d08..5c4f8dee2a10 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java @@ -211,13 +211,65 @@ record -> { */ public static BigQueryReadSettings.Builder createBigQueryReadSettingsBuilder() { BigQueryReadSettings.Builder builder = BigQueryReadSettings.newBuilder(); - String endpoint = System.getenv("BIGQUERY_STORAGE_ENDPOINT"); + String endpoint = + System.getProperty("bigquery.storage.endpoint", System.getenv("BIGQUERY_STORAGE_ENDPOINT")); if (endpoint != null) { builder.setEndpoint(endpoint); } return builder; } + /** Extracts the region name from the BigQuery endpoint, or returns null if not regional. */ + public static String getBigQueryRegion() { + String endpoint = System.getProperty("bigquery.endpoint", System.getenv("BIGQUERY_ENDPOINT")); + if (endpoint == null) { + return null; + } + if (endpoint.contains("-bigquery.googleapis.com")) { + int start = endpoint.indexOf("https://"); + start = (start == -1) ? 0 : start + 8; + int end = endpoint.indexOf("-bigquery.googleapis.com"); + return endpoint.substring(start, end); + } + if (endpoint.contains(".rep.googleapis.com") && endpoint.contains("bigquery.")) { + int start = endpoint.indexOf("bigquery.") + 9; + int end = endpoint.indexOf(".rep.googleapis.com"); + return endpoint.substring(start, end); + } + return null; + } + + /** + * Extracts the region name from the BigQuery Storage endpoint, or returns null if not regional. + */ + public static String getBigQueryStorageRegion() { + String endpoint = + System.getProperty("bigquery.storage.endpoint", System.getenv("BIGQUERY_STORAGE_ENDPOINT")); + if (endpoint == null) { + return null; + } + if (endpoint.contains("-bigquerystorage.googleapis.com")) { + int end = endpoint.indexOf("-bigquerystorage.googleapis.com"); + return endpoint.substring(0, end); + } + if (endpoint.contains(".rep.googleapis.com") && endpoint.contains("bigquerystorage.")) { + int start = endpoint.indexOf("bigquerystorage.") + 16; + int end = endpoint.indexOf(".rep.googleapis.com"); + return endpoint.substring(start, end); + } + return null; + } + + /** Helper to check if the BQ Storage client is configured to target a regional endpoint. */ + public static boolean isRegionalEndpoint() { + return getBigQueryStorageRegion() != null; + } + + /** Helper to check if the BigQuery API client is configured to target a regional endpoint. */ + public static boolean isBigQueryRegionalEndpoint() { + return getBigQueryRegion() != null; + } + /** * Returns a {@link BigQueryReadClient} configured with potential endpoint overrides for testing. */ @@ -231,7 +283,8 @@ public static BigQueryReadClient createBigQueryReadClient() throws IOException { */ public static BigQueryWriteSettings.Builder createBigQueryWriteSettingsBuilder() { BigQueryWriteSettings.Builder builder = BigQueryWriteSettings.newBuilder(); - String endpoint = System.getenv("BIGQUERY_STORAGE_ENDPOINT"); + String endpoint = + System.getProperty("bigquery.storage.endpoint", System.getenv("BIGQUERY_STORAGE_ENDPOINT")); if (endpoint != null) { builder.setEndpoint(endpoint); }