From e8d74fd60cc5190a87a141e17ca4d52e3c90360d Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 10 Aug 2026 13:20:26 -0400 Subject: [PATCH] test(bigquery): support regional endpoints dynamically in integration tests Configures the 'BIGQUERY_ENDPOINT' and 'BIGQUERY_STORAGE_ENDPOINT' environment variables in the Kokoro presubmit and GraalVM test configuration files to point to the 'us-east7' regional endpoints. Additionally, un-skips the integration tests when running against regional endpoints by instantiating separate global client instances inside the tests for the cases that query public datasets (which are only available globally). --- .../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 | 346 ++++++++++-------- .../bigquery/it/ITNightlyBigQueryTest.java | 16 +- .../bigquery/it/ITOpenTelemetryTest.java | 47 ++- .../it/ITBigQueryStorageReadClientTest.java | 313 +++++++++------- .../it/ITBigQueryStorageWriteClientTest.java | 34 +- .../bigquery/storage/v1/it/util/Helper.java | 57 ++- .../google-cloud-bigquerystorage/pom.xml | 6 + 13 files changed, 544 insertions(+), 326 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 2c5fc8df4e8d..49890d425a61 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..dddf2f88936b 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 @@ -1038,6 +1038,7 @@ class ITBigQueryTest { .build(); private static BigQuery bigquery; + private static BigQuery globalBigQuery; private static Storage storage; private static OpenTelemetry otel; @@ -1083,6 +1084,12 @@ static void beforeClass() throws InterruptedException, IOException { otel = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal(); bigquery = bigqueryHelper.getOptions().getService(); + globalBigQuery = + BigQueryOptions.newBuilder() + .setProjectId(bigqueryHelper.getOptions().getProjectId()) + .setCredentials(bigqueryHelper.getOptions().getCredentials()) + .build() + .getService(); storage = storageHelper.getOptions().getService(); storage.create(BucketInfo.of(BUCKET)); storage.create( @@ -1271,7 +1278,7 @@ void testLosslessMaxTimestampIntegration() throws InterruptedException { @Test void testListDatasets() { - Page datasets = bigquery.listDatasets("bigquery-public-data"); + Page datasets = globalBigQuery.listDatasets("bigquery-public-data"); Iterator iterator = datasets.iterateAll().iterator(); Set datasetNames = new HashSet<>(); Map datasetLocation = new HashMap<>(); @@ -2169,45 +2176,52 @@ void testCreateAndUpdateTableWithPolicyTags() throws IOException { Schema policySchema = Schema.of(STRING_FIELD_SCHEMA, stringFieldWithPolicy, INTEGER_FIELD_SCHEMA); - // Test: Amend an existing schema with a policy tag. - String tableNameForUpdate = "test_update_table_policytags"; - TableId tableIdForUpdate = TableId.of(DATASET, tableNameForUpdate); - TableInfo tableInfo = - TableInfo.newBuilder(tableIdForUpdate, StandardTableDefinition.of(TABLE_SCHEMA)) - .setDescription("policy tag update test table") - .build(); - Table createdTableForUpdate = bigquery.create(tableInfo); - assertNotNull(createdTableForUpdate); - Schema schema = createdTableForUpdate.getDefinition().getSchema(); - FieldList fields = schema.getFields(); - // Create a new schema adding the current fields, plus the new policy tag field - List fieldList = new ArrayList<>(); - for (Field field : fields) { - fieldList.add(field); + String policyDatasetName = "policyset_" + UUID.randomUUID().toString().replace("-", "_"); + globalBigQuery.create(DatasetInfo.newBuilder(policyDatasetName).setLocation("US").build()); + try { + // Test: Amend an existing schema with a policy tag. + String tableNameForUpdate = "test_update_table_policytags"; + TableId tableIdForUpdate = TableId.of(policyDatasetName, tableNameForUpdate); + TableInfo tableInfo = + TableInfo.newBuilder(tableIdForUpdate, StandardTableDefinition.of(TABLE_SCHEMA)) + .setDescription("policy tag update test table") + .build(); + Table createdTableForUpdate = globalBigQuery.create(tableInfo); + assertNotNull(createdTableForUpdate); + Schema schema = createdTableForUpdate.getDefinition().getSchema(); + FieldList fields = schema.getFields(); + // Create a new schema adding the current fields, plus the new policy tag field + List fieldList = new ArrayList<>(); + for (Field field : fields) { + fieldList.add(field); + } + fieldList.add(stringFieldWithPolicy); + Schema updatedSchemaWithPolicyTag = Schema.of(fieldList); + Table updatedTable = + createdTableForUpdate.toBuilder() + .setDefinition(StandardTableDefinition.of(updatedSchemaWithPolicyTag)) + .build(); + updatedTable.update(); + Table remoteUpdatedTable = globalBigQuery.getTable(policyDatasetName, tableNameForUpdate); + assertEquals( + updatedSchemaWithPolicyTag, + remoteUpdatedTable.getDefinition().getSchema()); + globalBigQuery.delete(tableIdForUpdate); + + // Test: Create a new table with a policy tag defined. + String tableName = "test_create_table_policytags"; + TableId tableId = TableId.of(policyDatasetName, tableName); + StandardTableDefinition tableDefinition = + StandardTableDefinition.newBuilder().setSchema(policySchema).build(); + Table createdTable = globalBigQuery.create(TableInfo.of(tableId, tableDefinition)); + assertNotNull(createdTable); + Table remoteTable = globalBigQuery.getTable(policyDatasetName, tableName); + assertEquals( + policySchema, remoteTable.getDefinition().getSchema()); + globalBigQuery.delete(tableId); + } finally { + RemoteBigQueryHelper.forceDelete(globalBigQuery, policyDatasetName); } - fieldList.add(stringFieldWithPolicy); - Schema updatedSchemaWithPolicyTag = Schema.of(fieldList); - Table updatedTable = - createdTableForUpdate.toBuilder() - .setDefinition(StandardTableDefinition.of(updatedSchemaWithPolicyTag)) - .build(); - updatedTable.update(); - Table remoteUpdatedTable = bigquery.getTable(DATASET, tableNameForUpdate); - assertEquals( - updatedSchemaWithPolicyTag, - remoteUpdatedTable.getDefinition().getSchema()); - bigquery.delete(tableIdForUpdate); - - // Test: Create a new table with a policy tag defined. - String tableName = "test_create_table_policytags"; - TableId tableId = TableId.of(DATASET, tableName); - StandardTableDefinition tableDefinition = - StandardTableDefinition.newBuilder().setSchema(policySchema).build(); - Table createdTable = bigquery.create(TableInfo.of(tableId, tableDefinition)); - assertNotNull(createdTable); - Table remoteTable = bigquery.getTable(DATASET, tableName); - assertEquals(policySchema, remoteTable.getDefinition().getSchema()); - bigquery.delete(tableId); // Clean up policy tags policyTagManagerClient.deletePolicyTag(policyTagId); @@ -2515,24 +2529,30 @@ void testCreateExternalTable() throws InterruptedException { @Test void testSetPermExternalTableSchema() { - String tableName = "test_create_external_table_perm"; - TableId tableId = TableId.of(DATASET, tableName); - ExternalTableDefinition externalTableDefinition = - ExternalTableDefinition.newBuilder( - "gs://" + BUCKET + "/" + JSON_LOAD_FILE, FormatOptions.json()) - .setSchema(TABLE_SCHEMA) - .setConnectionId( - "projects/java-docs-samples-testing/locations/us/connections/DEVREL_TEST_CONNECTION") - .build(); - TableInfo tableInfo = TableInfo.of(tableId, externalTableDefinition); - Table createdTable = bigquery.create(tableInfo); + String tempDataset = "extconnset_" + UUID.randomUUID().toString().replace("-", "_"); + globalBigQuery.create(DatasetInfo.newBuilder(tempDataset).setLocation("US").build()); + try { + String tableName = "test_create_external_table_perm"; + TableId tableId = TableId.of(tempDataset, tableName); + ExternalTableDefinition externalTableDefinition = + ExternalTableDefinition.newBuilder( + "gs://" + BUCKET + "/" + JSON_LOAD_FILE, FormatOptions.json()) + .setSchema(TABLE_SCHEMA) + .setConnectionId( + "projects/java-docs-samples-testing/locations/us/connections/DEVREL_TEST_CONNECTION") + .build(); + TableInfo tableInfo = TableInfo.of(tableId, externalTableDefinition); + Table createdTable = globalBigQuery.create(tableInfo); - assertNotNull(createdTable); - assertEquals(DATASET, createdTable.getTableId().getDataset()); - assertEquals(tableName, createdTable.getTableId().getTable()); - Table remoteTable = bigquery.getTable(DATASET, tableName); - assertNotNull(remoteTable); - assertTrue(remoteTable.delete()); + assertNotNull(createdTable); + assertEquals(tempDataset, createdTable.getTableId().getDataset()); + assertEquals(tableName, createdTable.getTableId().getTable()); + Table remoteTable = globalBigQuery.getTable(tempDataset, tableName); + assertNotNull(remoteTable); + assertTrue(remoteTable.delete()); + } finally { + RemoteBigQueryHelper.forceDelete(globalBigQuery, tempDataset); + } } @Test @@ -2931,11 +2951,11 @@ void testDeleteJob() { String jobName = "jobId_" + UUID.randomUUID().toString(); JobId jobId = JobId.newBuilder().setLocation("us-east1").setJob(jobName).setProject(PROJECT_ID).build(); - Job createdJob = bigquery.create(JobInfo.of(jobId, config)); - Job remoteJob = bigquery.getJob(createdJob.getJobId()); + Job createdJob = globalBigQuery.create(JobInfo.of(jobId, config)); + Job remoteJob = globalBigQuery.getJob(createdJob.getJobId()); assertEquals(createdJob.getJobId(), remoteJob.getJobId()); - assertTrue(bigquery.delete(jobId)); - assertNull(bigquery.getJob(jobId)); + assertTrue(globalBigQuery.delete(jobId)); + assertNull(globalBigQuery.getJob(jobId)); } @Test @@ -3190,10 +3210,10 @@ void testListAllTableData() { void testListPageWithStartIndex() { String tableName = "midyear_population_agespecific"; TableId tableId = TableId.of(PUBLIC_PROJECT, PUBLIC_DATASET, tableName); - Table table = bigquery.getTable(tableId); + Table table = globalBigQuery.getTable(tableId); long numRows = table.getNumRows().longValue(); Page tableResult = - bigquery.listTableData( + globalBigQuery.listTableData( tableId, BigQuery.TableDataListOption.startIndex(numRows - 300_000L), BigQuery.TableDataListOption.pageSize(600_000L)); @@ -3660,7 +3680,7 @@ void testQueryStatistics() throws InterruptedException { @Test void testExecuteSelectDefaultConnectionSettings() throws SQLException { // Use the default connection settings - Connection connection = bigquery.createConnection(); + Connection connection = globalBigQuery.createConnection(); String query = "SELECT corpus FROM `bigquery-public-data.samples.shakespeare` GROUP BY corpus;"; BigQueryResult bigQueryResult = connection.executeSelect(query); assertEquals(42, bigQueryResult.getTotalRows()); @@ -3672,7 +3692,7 @@ void testExecuteSelectWithReadApi() throws SQLException { final int rowLimit = 5000; final String QUERY = "SELECT * FROM bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2017 LIMIT %s"; - bigquery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED); + globalBigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED); // Set priority explicitly to ensure that fast query is not used. // min result size and page row count ratio ensure that the ReadAPI is used. ConnectionSettings connectionSettingsReadAPIEnabledFastQueryDisabled = @@ -3684,7 +3704,7 @@ void testExecuteSelectWithReadApi() throws SQLException { .build(); Connection connectionReadAPIEnabled = - bigquery.createConnection(connectionSettingsReadAPIEnabledFastQueryDisabled); + globalBigQuery.createConnection(connectionSettingsReadAPIEnabledFastQueryDisabled); String selectQuery = String.format(QUERY, rowLimit); @@ -3711,7 +3731,7 @@ void testExecuteSelectWithFastQueryReadApi() throws SQLException { .build(); Connection connectionReadAPIEnabled = - bigquery.createConnection(connectionSettingsReadAPIEnabledFastQueryDisabled); + globalBigQuery.createConnection(connectionSettingsReadAPIEnabledFastQueryDisabled); String selectQuery = String.format(QUERY, rowLimit); @@ -4792,11 +4812,11 @@ void testLocationFastSQLQueryWithJobId() throws InterruptedException { .setLocation("europe-west1") .setLabels(LABELS) .build(); - bigquery.create(infoUK); + globalBigQuery.create(infoUK); TableDefinition tableDefinition = StandardTableDefinition.of(SIMPLE_SCHEMA); TableInfo tableInfo = TableInfo.newBuilder(tableIdFastQueryUk, tableDefinition).build(); - bigquery.create(tableInfo); + globalBigQuery.create(tableInfo); String insert = "INSERT " + UK_DATASET + "." + tableIdFastQueryUk.getTable() + " VALUES('Anna');"; @@ -4805,11 +4825,11 @@ void testLocationFastSQLQueryWithJobId() throws InterruptedException { QueryJobConfiguration.newBuilder(insert) .setDefaultDataset(DatasetId.of(UK_DATASET)) .build(); - TableResult result = bigquery.query(config); + TableResult result = globalBigQuery.query(config); assertNotNull(result.getJobId()); assertEquals(SIMPLE_SCHEMA, result.getSchema()); // Use `getNumDmlAffectedRows()` for DML operations - Job queryJob = bigquery.getJob(result.getJobId()); + Job queryJob = globalBigQuery.getJob(result.getJobId()); queryJob = queryJob.waitFor(); assertNull( queryJob.getStatus().getError(), @@ -4834,7 +4854,7 @@ void testLocationFastSQLQueryWithJobId() throws InterruptedException { QueryJobConfiguration configSelect = QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(UK_DATASET)).build(); try { - bigquery.query(configSelect, jobIdWithLocation); + globalBigQuery.query(configSelect, jobIdWithLocation); } catch (BigQueryException exception) { assertTrue(exception.getMessage().contains("Not found")); assertEquals(BigQueryException.class, exception.getClass()); @@ -4844,7 +4864,7 @@ void testLocationFastSQLQueryWithJobId() throws InterruptedException { JobId jobIdNoLocation = JobId.newBuilder().build(); QueryJobConfiguration configNoLocation = QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(UK_DATASET)).build(); - TableResult resultNoLocation = bigquery.query(configNoLocation, jobIdNoLocation); + TableResult resultNoLocation = globalBigQuery.query(configNoLocation, jobIdNoLocation); for (FieldValueList row : resultNoLocation.getValues()) { FieldValue stringCell = row.get(0); assertEquals(stringCell, row.get("StringField")); @@ -4960,34 +4980,40 @@ void testFastDDLQuery() throws InterruptedException { @Test void testFastQuerySlowDDL() throws InterruptedException { - 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 = - String.format( - "CREATE OR REPLACE TABLE %s AS SELECT unique_key, agency, complaint_type, descriptor," - + " street_name, city, landmark FROM" - + " `bigquery-public-data.new_york.311_service_requests`", - tableName); - QueryJobConfiguration ddlConfig = - QueryJobConfiguration.newBuilder(slowDdlQuery) - .setDefaultDataset(DatasetId.of(DATASET)) - .build(); - TableResult result = bigquery.query(ddlConfig); - assertNotNull(result.getJobId()); - assertEquals(0, result.getTotalRows()); - assertNotNull(result.getSchema()); - // Verify correctness of table content - String sqlQuery = String.format("SELECT * FROM %s.%s", DATASET, tableName); - QueryJobConfiguration sqlConfig = QueryJobConfiguration.newBuilder(sqlQuery).build(); - TableResult resultAfterDDL = bigquery.query(sqlConfig); - assertNotNull(resultAfterDDL.getJobId()); - for (FieldValueList row : resultAfterDDL.getValues()) { - FieldValue unique_key = row.get(0); - assertEquals(unique_key, row.get("unique_key")); - FieldValue agency = row.get(1); - assertEquals(agency, row.get("agency")); - FieldValue complaint_type = row.get(2); - assertEquals(complaint_type, row.get("complaint_type")); + String tempDataset = "slowddlset_" + UUID.randomUUID().toString().replace("-", "_"); + globalBigQuery.create(DatasetInfo.newBuilder(tempDataset).setLocation("US").build()); + try { + 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 = + String.format( + "CREATE OR REPLACE TABLE %s AS SELECT unique_key, agency, complaint_type, descriptor," + + " street_name, city, landmark FROM" + + " `bigquery-public-data.new_york.311_service_requests`", + tableName); + QueryJobConfiguration ddlConfig = + QueryJobConfiguration.newBuilder(slowDdlQuery) + .setDefaultDataset(DatasetId.of(tempDataset)) + .build(); + TableResult result = globalBigQuery.query(ddlConfig); + assertNotNull(result.getJobId()); + assertEquals(0, result.getTotalRows()); + assertNotNull(result.getSchema()); + // Verify correctness of table content + String sqlQuery = String.format("SELECT * FROM %s.%s", tempDataset, tableName); + QueryJobConfiguration sqlConfig = QueryJobConfiguration.newBuilder(sqlQuery).build(); + TableResult resultAfterDDL = globalBigQuery.query(sqlConfig); + assertNotNull(resultAfterDDL.getJobId()); + for (FieldValueList row : resultAfterDDL.getValues()) { + FieldValue unique_key = row.get(0); + assertEquals(unique_key, row.get("unique_key")); + FieldValue agency = row.get(1); + assertEquals(agency, row.get("agency")); + FieldValue complaint_type = row.get(2); + assertEquals(complaint_type, row.get("complaint_type")); + } + } finally { + RemoteBigQueryHelper.forceDelete(globalBigQuery, tempDataset); } } @@ -5075,7 +5101,7 @@ void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedExcepti String sessionId; // Imports a local file into a table. - try (TableDataWriteChannel writer = bigquery.writer(jobId, configuration); + try (TableDataWriteChannel writer = globalBigQuery.writer(jobId, configuration); OutputStream stream = Channels.newOutputStream(writer)) { InputStream inputStream = ITBigQueryTest.class.getClassLoader().getResourceAsStream("sessionTest.csv"); @@ -5086,7 +5112,7 @@ void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedExcepti } catch (IOException e) { throw new RuntimeException(e); } - Job loadJob = bigquery.getJob(jobId); + Job loadJob = globalBigQuery.getJob(jobId); Job completedJob = loadJob.waitFor(); assertNotNull(completedJob); @@ -5109,7 +5135,7 @@ void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedExcepti .build(); String sessionJobName = "jobId_" + UUID.randomUUID().toString(); JobId sessionJobId = JobId.newBuilder().setLocation("us").setJob(sessionJobName).build(); - try (TableDataWriteChannel writer = bigquery.writer(sessionJobId, sessionConfiguration); + try (TableDataWriteChannel writer = globalBigQuery.writer(sessionJobId, sessionConfiguration); OutputStream stream = Channels.newOutputStream(writer)) { InputStream inputStream = ITBigQueryTest.class.getClassLoader().getResourceAsStream("sessionTest.csv"); @@ -5117,7 +5143,7 @@ void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedExcepti } catch (IOException e) { throw new RuntimeException(e); } - Job queryJobWithSession = bigquery.getJob(sessionJobId); + Job queryJobWithSession = globalBigQuery.getJob(sessionJobId); queryJobWithSession = queryJobWithSession.waitFor(); LoadStatistics statisticsWithSession = queryJobWithSession.getStatistics(); assertNotNull(statisticsWithSession.getSessionInfo().getSessionId()); @@ -5128,7 +5154,7 @@ void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedExcepti QueryJobConfiguration.newBuilder(queryTempTable) .setConnectionProperties(ImmutableList.of(sessionConnectionProperty)) .build(); - Job queryTempTableJob = bigquery.create(JobInfo.of(queryJobConfigurationWithSession)); + Job queryTempTableJob = globalBigQuery.create(JobInfo.of(queryJobConfigurationWithSession)); queryTempTableJob = queryTempTableJob.waitFor(); assertNotNull(queryTempTableJob.getQueryResults()); } @@ -5306,12 +5332,12 @@ void testScriptStatistics() throws InterruptedException { + " FROM `bigquery-public-data`.samples.shakespeare\n" + ");"; QueryJobConfiguration config = QueryJobConfiguration.of(script); - Job remoteJob = bigquery.create(JobInfo.of(config)); + Job remoteJob = globalBigQuery.create(JobInfo.of(config)); JobInfo info = remoteJob.waitFor(); JobStatistics jobStatistics = info.getStatistics(); String parentJobId = info.getJobId().getJob(); assertEquals(2, jobStatistics.getNumChildJobs().longValue()); - Page page = bigquery.listJobs(JobListOption.parentJobId(parentJobId)); + Page page = globalBigQuery.listJobs(JobListOption.parentJobId(parentJobId)); for (Job job : page.iterateAll()) { JobStatistics.ScriptStatistics scriptStatistics = job.getStatistics().getScriptStatistics(); if (scriptStatistics != null) { @@ -6541,7 +6567,7 @@ void testCancelJob() throws InterruptedException, TimeoutException { @Test void testCancelNonExistingJob() { - assertFalse(bigquery.cancel("test_cancel_non_existing_job")); + assertFalse(globalBigQuery.cancel("test_cancel_non_existing_job")); } @Test @@ -6684,7 +6710,7 @@ void testLocation() throws Exception { Tracer tracer = otel.getTracer("Test Tracer"); BigQuery otelBigquery = - bigquery.getOptions().toBuilder() + globalBigQuery.getOptions().toBuilder() .setEnableOpenTelemetryTracing(true) .setOpenTelemetryTracer(tracer) .build() @@ -6692,12 +6718,12 @@ void testLocation() throws Exception { String datasetName = "locationset_" + UUID.randomUUID().toString().replace("-", "_"); Dataset dataset = - otelBigquery.create(DatasetInfo.newBuilder(datasetName).setLocation(location).build()); + globalBigQuery.create(DatasetInfo.newBuilder(datasetName).setLocation(location).build()); try { TableId tableId = TableId.of(dataset.getDatasetId().getDataset(), "sometable"); Schema schema = Schema.of(Field.of("name", LegacySQLTypeName.STRING)); TableDefinition tableDef = StandardTableDefinition.of(schema); - Table table = otelBigquery.create(TableInfo.newBuilder(tableId, tableDef).build()); + Table table = globalBigQuery.create(TableInfo.newBuilder(tableId, tableDef).build()); String query = String.format( @@ -6777,7 +6803,7 @@ void testLocation() throws Exception { "writing to a table with wrong location shouldn't work"); } } finally { - RemoteBigQueryHelper.forceDelete(bigquery, datasetName); + RemoteBigQueryHelper.forceDelete(globalBigQuery, datasetName); } } @@ -7448,9 +7474,8 @@ void testStatelessQueriesWithLocation() throws Exception { String location = "EU"; String wrongLocation = "US"; - RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); BigQuery bigQuery = - bigqueryHelper.getOptions().toBuilder().setLocation(location).build().getService(); + globalBigQuery.getOptions().toBuilder().setLocation(location).build().getService(); String datasetName = "locationset_" + UUID.randomUUID().toString().replace("-", "_"); Dataset dataset = @@ -7478,7 +7503,7 @@ void testStatelessQueriesWithLocation() throws Exception { BigQueryException.class, () -> { BigQuery bigQueryWrongLocation = - bigqueryHelper.getOptions().toBuilder() + globalBigQuery.getOptions().toBuilder() .setLocation(wrongLocation) .build() .getService(); @@ -7609,9 +7634,8 @@ void testInvalidUniverseDomainWithMismatchCredentials() { @Test void testUniverseDomainWithMatchingDomain() { // Test a valid domain using the default credentials and Google default universe domain. - RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); BigQueryOptions bigQueryOptions = - bigqueryHelper.getOptions().toBuilder().setUniverseDomain("googleapis.com").build(); + globalBigQuery.getOptions().toBuilder().setUniverseDomain("googleapis.com").build(); BigQuery bigQuery = bigQueryOptions.getService(); // Verify that all is well by listing a dataset. @@ -7701,44 +7725,50 @@ void testExternalMetadataCacheModeFailForNonBiglake() { @Test void testObjectTable() throws InterruptedException { - String tableName = generateRandomName("test_object_table"); - TableId tableId = TableId.of(DATASET, tableName); - - String sourceUri = "gs://" + BUCKET + "/" + JSON_LOAD_FILE; - ExternalTableDefinition externalTableDefinition = - ExternalTableDefinition.newBuilder(sourceUri) - .setConnectionId( - "projects/java-docs-samples-testing/locations/us/connections/DEVREL_TEST_CONNECTION") - .setObjectMetadata("SIMPLE") - .build(); - TableInfo tableInfo = TableInfo.of(tableId, externalTableDefinition); - Table createdTable = bigquery.create(tableInfo); - assertNotNull(createdTable); - assertEquals(DATASET, createdTable.getTableId().getDataset()); - assertEquals(tableName, createdTable.getTableId().getTable()); - Table remoteTable = bigquery.getTable(DATASET, tableName); - assertNotNull(remoteTable); - + String tempDataset = "objectset_" + UUID.randomUUID().toString().replace("-", "_"); + globalBigQuery.create(DatasetInfo.newBuilder(tempDataset).setLocation("US").build()); try { - assertTrue(remoteTable.getDefinition() instanceof ExternalTableDefinition); - assertEquals(createdTable.getTableId(), remoteTable.getTableId()); - assertEquals( - "SIMPLE", ((ExternalTableDefinition) remoteTable.getDefinition()).getObjectMetadata()); - assertNotNull(remoteTable.getDefinition().getSchema().getFields().get("uri")); - - String query = String.format("SELECT * FROM %s.%s", DATASET, tableName); - QueryJobConfiguration config = QueryJobConfiguration.newBuilder(query).build(); - - Job remoteJob = bigquery.create(JobInfo.of(config)); - remoteJob = remoteJob.waitFor(); - assertNull(remoteJob.getStatus().getError()); + String tableName = generateRandomName("test_object_table"); + TableId tableId = TableId.of(tempDataset, tableName); + + String sourceUri = "gs://" + BUCKET + "/" + JSON_LOAD_FILE; + ExternalTableDefinition externalTableDefinition = + ExternalTableDefinition.newBuilder(sourceUri) + .setConnectionId( + "projects/java-docs-samples-testing/locations/us/connections/DEVREL_TEST_CONNECTION") + .setObjectMetadata("SIMPLE") + .build(); + TableInfo tableInfo = TableInfo.of(tableId, externalTableDefinition); + Table createdTable = globalBigQuery.create(tableInfo); + assertNotNull(createdTable); + assertEquals(tempDataset, createdTable.getTableId().getDataset()); + assertEquals(tableName, createdTable.getTableId().getTable()); + Table remoteTable = globalBigQuery.getTable(tempDataset, tableName); + assertNotNull(remoteTable); - Job queryJob = bigquery.getJob(remoteJob.getJobId()); - JobStatistics.QueryStatistics statistics = queryJob.getStatistics(); - assertNotNull(statistics); - assertThat(statistics.getTotalBytesProcessed()).isGreaterThan(0); + try { + assertTrue(remoteTable.getDefinition() instanceof ExternalTableDefinition); + assertEquals(createdTable.getTableId(), remoteTable.getTableId()); + assertEquals( + "SIMPLE", ((ExternalTableDefinition) remoteTable.getDefinition()).getObjectMetadata()); + assertNotNull(remoteTable.getDefinition().getSchema().getFields().get("uri")); + + String query = String.format("SELECT * FROM %s.%s", tempDataset, tableName); + QueryJobConfiguration config = QueryJobConfiguration.newBuilder(query).build(); + + Job remoteJob = globalBigQuery.create(JobInfo.of(config)); + remoteJob = remoteJob.waitFor(); + assertNull(remoteJob.getStatus().getError()); + + Job queryJob = globalBigQuery.getJob(remoteJob.getJobId()); + JobStatistics.QueryStatistics statistics = queryJob.getStatistics(); + assertNotNull(statistics); + assertThat(statistics.getTotalBytesProcessed()).isGreaterThan(0); + } finally { + assertTrue(remoteTable.delete()); + } } finally { - assertTrue(remoteTable.delete()); + RemoteBigQueryHelper.forceDelete(globalBigQuery, tempDataset); } } @@ -7848,12 +7878,12 @@ void testStatementType() throws InterruptedException { public void testOpenTelemetryTracingDatasets() { Tracer tracer = otel.getTracer("Test Tracer"); BigQueryOptions otelOptions = - BigQueryOptions.newBuilder() + globalBigQuery.getOptions().toBuilder() .setEnableOpenTelemetryTracing(true) .setOpenTelemetryTracer(tracer) .setLocation("US") .build(); - BigQuery bigquery = otelOptions.getService(); + BigQuery otelBigquery = otelOptions.getService(); Span parentSpan = tracer @@ -7872,9 +7902,9 @@ public void testOpenTelemetryTracingDatasets() { .setLocation("US") .build(); - Dataset dataset = bigquery.create(info); + Dataset dataset = otelBigquery.create(info); assertNotNull(dataset); - dataset = bigquery.getDataset(dataset.getDatasetId().getDataset()); + dataset = otelBigquery.getDataset(dataset.getDatasetId().getDataset()); assertNotNull(dataset); DatasetInfo updatedInfo = @@ -7884,9 +7914,9 @@ public void testOpenTelemetryTracingDatasets() { .setLabels(LABELS) .build(); - dataset = bigquery.update(updatedInfo, DatasetOption.accessPolicyVersion(2)); + dataset = otelBigquery.update(updatedInfo, DatasetOption.accessPolicyVersion(2)); assertEquals("Updated Description", dataset.getDescription()); - assertTrue(bigquery.delete(dataset.getDatasetId())); + assertTrue(otelBigquery.delete(dataset.getDatasetId())); } finally { parentSpan.end(); Map, Object> createMap = 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..2a825847c23b 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 @@ -89,6 +89,7 @@ public class ITNightlyBigQueryTest { 300000; // Used for multiquery testcase, a lower limit like 300K should be fine private static int rowCnt = 0; private static BigQuery bigquery; + private static BigQuery globalBigQuery; private static final String BASE_QUERY = "select StringField, GeographyField, BooleanField, BigNumericField, IntegerField, NumericField, BytesField, " + "TimestampField, TimeField, DateField, IntegerArrayField, RecordField.BooleanField, RecordField.StringField ," @@ -177,6 +178,12 @@ public class ITNightlyBigQueryTest { public static void beforeClass() throws InterruptedException, IOException { RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); bigquery = bigqueryHelper.getOptions().getService(); + globalBigQuery = + BigQueryOptions.newBuilder() + .setProjectId(bigqueryHelper.getOptions().getProjectId()) + .setCredentials(bigqueryHelper.getOptions().getCredentials()) + .build() + .getService(); createDataset(DATASET); createTable(DATASET, TABLE, BQ_SCHEMA); populateTestRecords(DATASET, TABLE); @@ -514,7 +521,7 @@ void testForTableNotFound() throws SQLException { logger.log(Level.INFO, "Query used: {0}", query); String dataSet = RemoteBigQueryHelper.generateDatasetName(); String table = "TAB_" + UUID.randomUUID(); - createDataset(dataSet); + globalBigQuery.create(DatasetInfo.newBuilder(dataSet).build()); TableId targetTable = TableId.of( ServiceOptions.getDefaultProjectId(), @@ -528,8 +535,7 @@ void testForTableNotFound() throws SQLException { .setAllowLargeResults(true) .build(); - Connection connection = - BigQueryOptions.getDefaultInstance().getService().createConnection(conSet); + Connection connection = globalBigQuery.createConnection(conSet); BigQueryResult bigQueryResultSet = connection.executeSelect(query); assertNotNull(getResultHashWiki(bigQueryResultSet)); // this iterated through all the rows assertTrue( @@ -539,8 +545,8 @@ void testForTableNotFound() throws SQLException { .getTotalRows())); // either job should return the actual count or -1 if the job // is still running try { - deleteTable(dataSet, table); - deleteDataset(dataSet); + globalBigQuery.delete(TableId.of(dataSet, table)); + globalBigQuery.delete(DatasetId.of(dataSet)); } catch (Exception e) { logger.log( Level.WARNING, 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-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java b/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java index 6accfa6bd878..81d826ea735a 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java @@ -132,6 +132,7 @@ class ITBigQueryStorageReadClientTest { private static final int MAX_STREAM_COUNT = 1; private static BigQueryReadClient readClient; + private static BigQueryReadClient globalReadClient; private static String projectName; private static String parentProjectId; private static BigQuery bigquery; @@ -504,7 +505,8 @@ 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(); + globalReadClient = Helper.isRegionalEndpoint() ? BigQueryReadClient.create() : readClient; projectName = ServiceOptions.getDefaultProjectId(); parentProjectId = String.format("projects/%s", projectName); @@ -515,10 +517,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); @@ -588,6 +599,11 @@ static void afterAll() throws InterruptedException { readClient.awaitTermination(10, TimeUnit.SECONDS); } + if (globalReadClient != null && globalReadClient != readClient) { + globalReadClient.close(); + globalReadClient.awaitTermination(10, TimeUnit.SECONDS); + } + if (bigquery != null) { RemoteBigQueryHelper.forceDelete(bigquery, DATASET); LOG.info("Deleted test dataset: " + DATASET); @@ -603,7 +619,7 @@ void testSimpleReadAvro() { /* tableId= */ "shakespeare"); ReadSession session = - readClient.createReadSession( + globalReadClient.createReadSession( /* parent= */ parentProjectId, /* readSession= */ ReadSession.newBuilder() .setTable(table) @@ -622,7 +638,8 @@ void testSimpleReadAvro() { ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build(); long rowCount = 0; - ServerStream stream = readClient.readRowsCallable().call(readRowsRequest); + ServerStream stream = + globalReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { rowCount += response.getRowCount(); } @@ -639,7 +656,7 @@ void testSimpleReadArrow() { /* tableId= */ "shakespeare"); ReadSession session = - readClient.createReadSession( + globalReadClient.createReadSession( /* parent= */ parentProjectId, /* readSession= */ ReadSession.newBuilder() .setTable(table) @@ -668,7 +685,8 @@ void testSimpleReadArrow() { long rowCount = 0; // Process each block of rows as they arrive and decode using our simple row reader. - ServerStream stream = readClient.readRowsCallable().call(readRowsRequest); + ServerStream stream = + globalReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { Preconditions.checkState(response.hasArrowRecordBatch()); rowCount += response.getRowCount(); @@ -947,7 +965,7 @@ void timestamp_readAvro() throws IOException { } @Test - void testSimpleReadAndResume() { + void testSimpleReadAndResume() throws IOException { String table = BigQueryResource.formatTableResource( /* projectId= */ "bigquery-public-data", @@ -955,7 +973,7 @@ void testSimpleReadAndResume() { /* tableId= */ "shakespeare"); ReadSession session = - readClient.createReadSession( + globalReadClient.createReadSession( /* parent= */ parentProjectId, /* readSession= */ ReadSession.newBuilder() .setTable(table) @@ -970,7 +988,8 @@ void testSimpleReadAndResume() { + " response:%n%s", table, session.toString())); - long rowCount = readStreamToOffset(session.getStreams(0), /* rowOffset= */ 34_846); + long rowCount = + readStreamToOffset(globalReadClient, session.getStreams(0), /* rowOffset= */ 34_846); ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder() @@ -978,7 +997,8 @@ void testSimpleReadAndResume() { .setOffset(rowCount) .build(); - ServerStream stream = readClient.readRowsCallable().call(readRowsRequest); + ServerStream stream = + globalReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { rowCount += response.getRowCount(); @@ -1012,7 +1032,7 @@ void testFilter() throws IOException { .build()) .build(); - ReadSession session = readClient.createReadSession(request); + ReadSession session = globalReadClient.createReadSession(request); assertEquals( 1, session.getStreamsCount(), @@ -1029,7 +1049,8 @@ void testFilter() throws IOException { long rowCount = 0; - ServerStream stream = readClient.readRowsCallable().call(readRowsRequest); + ServerStream stream = + globalReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { rowCount += response.getRowCount(); reader.processRows( @@ -1077,7 +1098,7 @@ void testColumnSelection() throws IOException { .build()) .build(); - ReadSession session = readClient.createReadSession(request); + ReadSession session = globalReadClient.createReadSession(request); assertEquals( 1, session.getStreamsCount(), @@ -1108,7 +1129,8 @@ void testColumnSelection() throws IOException { SimpleRowReaderAvro reader = new SimpleRowReaderAvro(avroSchema); long rowCount = 0; - ServerStream stream = readClient.readRowsCallable().call(readRowsRequest); + ServerStream stream = + globalReadClient.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : stream) { rowCount += response.getRowCount(); reader.processRows( @@ -1597,53 +1619,58 @@ void testStructAndArraySqlTypes() throws InterruptedException, IOException { @Test void testSimpleReadWithBackgroundExecutorProvider() throws IOException { - BigQueryReadSettings bigQueryReadSettings = - com.google.cloud.bigquery.storage.v1.it.util.Helper.createBigQueryReadSettingsBuilder() + BigQueryReadSettings.Builder settingsBuilder = + Helper.createBigQueryReadSettingsBuilder() .setBackgroundExecutorProvider( - InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(14).build()) - .build(); - // Overriding the default client - readClient = BigQueryReadClient.create(bigQueryReadSettings); - assertTrue( - readClient.getStub().getStubSettings().getBackgroundExecutorProvider() - instanceof InstantiatingExecutorProvider); - assertEquals( - 14, - ((InstantiatingExecutorProvider) - readClient.getStub().getStubSettings().getBackgroundExecutorProvider()) - .getExecutorThreadCount()); - String table = - BigQueryResource.formatTableResource( - /* projectId= */ "bigquery-public-data", - /* datasetId= */ "samples", - /* tableId= */ "shakespeare"); + InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(14).build()); + if (Helper.isRegionalEndpoint()) { + settingsBuilder.setEndpoint(BigQueryReadSettings.getDefaultEndpoint()); + } + BigQueryReadClient localClient = BigQueryReadClient.create(settingsBuilder.build()); + try { + assertTrue( + localClient.getStub().getStubSettings().getBackgroundExecutorProvider() + instanceof InstantiatingExecutorProvider); + assertEquals( + 14, + ((InstantiatingExecutorProvider) + localClient.getStub().getStubSettings().getBackgroundExecutorProvider()) + .getExecutorThreadCount()); + 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); + assertEquals( + 1, + session.getStreamsCount(), + String.format( + "Did not receive expected number of streams for table '%s' CreateReadSession" + + " response:%n%s", + table, session.toString())); - ReadSession session = - readClient.createReadSession( - /* parent= */ parentProjectId, - /* readSession= */ ReadSession.newBuilder() - .setTable(table) - .setDataFormat(DataFormat.AVRO) - .build(), - /* maxStreamCount= */ 1); - assertEquals( - 1, - session.getStreamsCount(), - String.format( - "Did not receive expected number of streams for table '%s' CreateReadSession" - + " response:%n%s", - table, session.toString())); + ReadRowsRequest readRowsRequest = + ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build(); - ReadRowsRequest readRowsRequest = - ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build(); + long rowCount = 0; + ServerStream stream = localClient.readRowsCallable().call(readRowsRequest); + for (ReadRowsResponse response : stream) { + rowCount += response.getRowCount(); + } - long rowCount = 0; - ServerStream stream = readClient.readRowsCallable().call(readRowsRequest); - for (ReadRowsResponse response : stream) { - rowCount += response.getRowCount(); + assertEquals(SHAKESPEARE_SAMPLE_ROW_COUNT, rowCount); + } finally { + localClient.close(); } - - assertEquals(SHAKESPEARE_SAMPLE_ROW_COUNT, rowCount); } @Test @@ -1722,33 +1749,35 @@ void testUniverseDomainWithMatchingDomain() throws IOException { BigQueryReadSettings bigQueryReadSettings = BigQueryReadSettings.newBuilder().setUniverseDomain("googleapis.com").build(); BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings); + try { + 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); - 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(); - ReadRowsRequest readRowsRequest = - ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build(); + long rowCount = 0; + ServerStream stream = localClient.readRowsCallable().call(readRowsRequest); + for (ReadRowsResponse response : stream) { + rowCount += response.getRowCount(); + } - long rowCount = 0; - ServerStream stream = readClient.readRowsCallable().call(readRowsRequest); - for (ReadRowsResponse response : stream) { - rowCount += response.getRowCount(); + assertEquals(SHAKESPEARE_SAMPLE_ROW_COUNT, rowCount); + } finally { + localClient.close(); } - - assertEquals(SHAKESPEARE_SAMPLE_ROW_COUNT, rowCount); - localClient.close(); } @Test @@ -1760,62 +1789,68 @@ void testSimpleReadWithOtelTracing() throws IOException { .build(); OpenTelemetry otel = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); - BigQueryReadSettings otelSettings = - com.google.cloud.bigquery.storage.v1.it.util.Helper.createBigQueryReadSettingsBuilder() + BigQueryReadSettings.Builder otelSettingsBuilder = + Helper.createBigQueryReadSettingsBuilder() .setEnableOpenTelemetryTracing(true) - .setOpenTelemetryTracerProvider(tracerProvider) - .build(); - BigQueryReadClient otelClient = BigQueryReadClient.create(otelSettings); - - String table = - BigQueryResource.formatTableResource( - /* projectId= */ "bigquery-public-data", - /* datasetId= */ "samples", - /* tableId= */ "shakespeare"); - - ReadSession session = - otelClient.createReadSession( - /* parent= */ parentProjectId, - /* readSession= */ ReadSession.newBuilder() - .setTable(table) - .setDataFormat(DataFormat.AVRO) - .build(), - /* maxStreamCount= */ 1); - - ReadRowsRequest readRowsRequest = - ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build(); + .setOpenTelemetryTracerProvider(tracerProvider); + if (Helper.isRegionalEndpoint()) { + otelSettingsBuilder.setEndpoint(BigQueryReadSettings.getDefaultEndpoint()); + } + BigQueryReadClient otelClient = BigQueryReadClient.create(otelSettingsBuilder.build()); + try { + String table = + BigQueryResource.formatTableResource( + /* projectId= */ "bigquery-public-data", + /* datasetId= */ "samples", + /* tableId= */ "shakespeare"); + + ReadSession session = + otelClient.createReadSession( + /* parent= */ parentProjectId, + /* readSession= */ ReadSession.newBuilder() + .setTable(table) + .setDataFormat(DataFormat.AVRO) + .build(), + /* maxStreamCount= */ 1); - ServerStream stream = otelClient.readRowsCallable().call(readRowsRequest); - - assertNotNull( - OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.createReadSession")); - assertNotNull( - OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.createReadSessionCallable")); - assertNotNull( - OTEL_ATTRIBUTES.get( - "com.google.cloud.bigquery.storage.v1.read.stub.createReadSessionCallable")); - assertNotNull( - OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.readRowsCallable")); - assertNotNull( - OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.stub.readRowsCallable")); - - // createReadSession is the parent span of createReadSessionCallable - assertEquals( - "com.google.cloud.bigquery.storage.v1.read.createReadSession", - OTEL_SPAN_IDS_TO_NAMES.get( - OTEL_PARENT_SPAN_IDS.get( - "com.google.cloud.bigquery.storage.v1.read.createReadSessionCallable"))); - - Map, Object> createReadSessionMap = - OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.createReadSession"); - assertNotNull(createReadSessionMap); - assertNotNull( - createReadSessionMap.get( - AttributeKey.longKey("bq.storage.read_session.request.max_stream_count"))); - assertEquals( - 1L, - createReadSessionMap.get( - AttributeKey.longKey("bq.storage.read_session.request.max_stream_count"))); + ReadRowsRequest readRowsRequest = + ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build(); + + ServerStream stream = otelClient.readRowsCallable().call(readRowsRequest); + + assertNotNull( + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.createReadSession")); + assertNotNull( + OTEL_ATTRIBUTES.get( + "com.google.cloud.bigquery.storage.v1.read.createReadSessionCallable")); + assertNotNull( + OTEL_ATTRIBUTES.get( + "com.google.cloud.bigquery.storage.v1.read.stub.createReadSessionCallable")); + assertNotNull( + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.readRowsCallable")); + assertNotNull( + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.stub.readRowsCallable")); + + // createReadSession is the parent span of createReadSessionCallable + assertEquals( + "com.google.cloud.bigquery.storage.v1.read.createReadSession", + OTEL_SPAN_IDS_TO_NAMES.get( + OTEL_PARENT_SPAN_IDS.get( + "com.google.cloud.bigquery.storage.v1.read.createReadSessionCallable"))); + + Map, Object> createReadSessionMap = + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.storage.v1.read.createReadSession"); + assertNotNull(createReadSessionMap); + assertNotNull( + createReadSessionMap.get( + AttributeKey.longKey("bq.storage.read_session.request.max_stream_count"))); + assertEquals( + 1L, + createReadSessionMap.get( + AttributeKey.longKey("bq.storage.read_session.request.max_stream_count"))); + } finally { + otelClient.close(); + } } void testUniverseDomain() throws IOException { @@ -1862,13 +1897,13 @@ void testUniverseDomain() throws IOException { * @param rowOffset * @return the number of requested rows to skip or the total rows read if stream had less rows. */ - private long readStreamToOffset(ReadStream readStream, long rowOffset) { + private long readStreamToOffset( + BigQueryReadClient client, ReadStream readStream, long rowOffset) { ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder().setReadStream(readStream.getName()).build(); long rowCount = 0; - ServerStream serverStream = - readClient.readRowsCallable().call(readRowsRequest); + ServerStream serverStream = client.readRowsCallable().call(readRowsRequest); for (ReadRowsResponse response : serverStream) { rowCount += response.getRowCount(); diff --git a/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java b/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java index 2e80cdb443b6..155969c1d211 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java @@ -114,6 +114,7 @@ class ITBigQueryStorageWriteClientTest { private static BigQueryReadClient readClient; private static BigQueryWriteClient writeClient; + private static BigQueryWriteClient globalWriteClient; private static String parentProjectId; private static TableInfo tableInfo; private static TableInfo tableInfo2; @@ -187,19 +188,29 @@ 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); + globalWriteClient = Helper.isRegionalEndpoint() ? BigQueryWriteClient.create() : writeClient; parentProjectId = String.format("projects/%s", ServiceOptions.getDefaultProjectId()); 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 = @@ -288,6 +299,11 @@ static void afterAll() throws InterruptedException { writeClient.awaitTermination(10, TimeUnit.SECONDS); } + if (globalWriteClient != null && globalWriteClient != writeClient) { + globalWriteClient.close(); + globalWriteClient.awaitTermination(10, TimeUnit.SECONDS); + } + if (readClient != null) { readClient.close(); readClient.awaitTermination(10, TimeUnit.SECONDS); @@ -374,7 +390,7 @@ ProtoRows createProtoRowsMixed(StringWithSecondsNanos[] messages) { void testBatchWriteWithCommittedStreamEU() throws IOException, InterruptedException, ExecutionException { WriteStream writeStream = - writeClient.createWriteStream( + globalWriteClient.createWriteStream( CreateWriteStreamRequest.newBuilder() .setParent(tableIdEU) .setWriteStream( @@ -383,7 +399,7 @@ void testBatchWriteWithCommittedStreamEU() ApiFuture response1; ApiFuture response2; try (StreamWriter streamWriter = - StreamWriter.newBuilder(writeStream.getName()) + StreamWriter.newBuilder(writeStream.getName(), globalWriteClient) .setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor())) .build()) { LOG.info("Sending one message"); @@ -2253,8 +2269,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-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java b/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java index 8e66820b6d08..5c4f8dee2a10 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage-it/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage-it/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); } diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml index 8cb32e3541e6..8a5723c22614 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml +++ b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml @@ -15,6 +15,8 @@ google-cloud-bigquerystorage + + @@ -408,6 +410,10 @@ --no-fallback --no-server + + ${bigquery.endpoint} + ${bigquery.storage.endpoint} +