From e4d359ca0fb7dad9913ba2620b13584d646e06ec Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Wed, 17 Jun 2026 13:02:06 -0400 Subject: [PATCH 1/2] fix(bigquery-jdbc): explicitly set location when available for temporary dataset --- .../bigquery/jdbc/BigQueryStatement.java | 12 +++-- .../bigquery/jdbc/BigQueryStatementTest.java | 44 +++++++++++++++++++ .../bigquery/jdbc/it/ITStatementTest.java | 37 ++++++++++++++++ 3 files changed, 89 insertions(+), 4 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 0d4d94175b8d..963ac6db35df 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -1375,11 +1375,15 @@ private void checkIfDatasetExistElseCreate(String datasetName) { Dataset dataset = bigQuery.getDataset(DatasetId.of(datasetName)); if (dataset == null) { LOG.info("Creating a hidden dataset: %s ", datasetName); - DatasetInfo datasetInfo = + DatasetInfo.Builder datasetInfoBuilder = DatasetInfo.newBuilder(datasetName) - .setDefaultTableLifetime(this.querySettings.getDestinationDatasetExpirationTime()) - .build(); - bigQuery.create(datasetInfo); + .setDefaultTableLifetime(this.querySettings.getDestinationDatasetExpirationTime()); + if (this.connection != null + && this.connection.getLocation() != null + && !this.connection.getLocation().isEmpty()) { + datasetInfoBuilder.setLocation(this.connection.getLocation()); + } + bigQuery.create(datasetInfoBuilder.build()); } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 674eb0df64e0..cf2861412cfd 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -34,6 +34,8 @@ import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.QueryResultsOption; import com.google.cloud.bigquery.BigQueryOptions; +import com.google.cloud.bigquery.DatasetId; +import com.google.cloud.bigquery.DatasetInfo; import com.google.cloud.bigquery.Field; import com.google.cloud.bigquery.FieldList; import com.google.cloud.bigquery.FieldValueList; @@ -707,4 +709,46 @@ public void testSetAndGetFetchSize() throws SQLException { bigQueryStatement.setFetchSize(100); assertEquals(100, bigQueryStatement.getFetchSize()); } + + @Test + public void testTemporaryDatasetCreationRespectsConnectionLocation() + throws SQLException, InterruptedException { + // 1. Setup mock connection with location and destination dataset + doReturn("europe-west3").when(bigQueryConnection).getLocation(); + doReturn("temp_dataset").when(bigQueryConnection).getDestinationDataset(); + + // Re-instantiate bigQueryStatement so it regenerates querySettings with these mocked connection + // values + BigQueryStatement statement = new BigQueryStatement(bigQueryConnection); + BigQueryStatement statementSpy = Mockito.spy(statement); + + // 2. Mock bigQuery.getDataset to return null (triggering creation) + doReturn(null).when(bigquery).getDataset(eq(DatasetId.of("temp_dataset"))); + + // 2b. Mock bigQuery.create for dry run during getStatementType + Job dryRunJobMock = getJobMock(null, null, StatementType.SELECT); + doReturn(dryRunJobMock).when(bigquery).create(any(JobInfo.class)); + + // 3. Mock bigquery.queryWithTimeout(...) to return tableResult (so execution doesn't fail on + // query execution) + TableResult result = mock(TableResult.class); + doReturn(result) + .when(bigquery) + .queryWithTimeout(any(QueryJobConfiguration.class), any(JobId.class), any()); + doReturn(mock(BigQueryJsonResultSet.class)) + .when(statementSpy) + .processJsonResultSet(eq(result), any()); + + // 4. Capture DatasetInfo passed to bigQuery.create() + ArgumentCaptor datasetInfoCaptor = ArgumentCaptor.forClass(DatasetInfo.class); + + // 5. Execute query + statementSpy.executeQuery("SELECT 1"); + + // 6. Verify dataset was created with correct location + verify(bigquery).create(datasetInfoCaptor.capture()); + DatasetInfo createdDatasetInfo = datasetInfoCaptor.getValue(); + assertEquals("temp_dataset", createdDatasetInfo.getDatasetId().getDataset()); + assertEquals("europe-west3", createdDatasetInfo.getLocation()); + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java index bc3a995d837a..2aaa750b3ea0 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java @@ -25,6 +25,10 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.cloud.ServiceOptions; +import com.google.cloud.bigquery.BigQuery; +import com.google.cloud.bigquery.BigQueryOptions; +import com.google.cloud.bigquery.Dataset; +import com.google.cloud.bigquery.DatasetId; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; @@ -412,4 +416,37 @@ int getSizeOfResultSet(ResultSet resultSet) throws SQLException { } return count; } + + @Test + public void testTemporaryDatasetLocation() throws SQLException, InterruptedException { + String projectId = DEFAULT_CATALOG; + String location = "europe-west3"; + String randomSuffix = String.valueOf(100 + new Random().nextInt(900)); + String tempDatasetName = "jdbc_temp_dataset_" + System.currentTimeMillis() + "_" + randomSuffix; + + String customConnectionUrl = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;ProjectId=" + + projectId + + ";OAuthType=3;Timeout=3600;Location=" + + location + + ";AllowLargeResults=true;LargeResultsDataset=" + + tempDatasetName + + ";"; + + BigQuery bigQuery = BigQueryOptions.getDefaultInstance().getService(); + try (Connection connection = DriverManager.getConnection(customConnectionUrl)) { + Statement statement = connection.createStatement(); + String query = "SELECT 1 as val;"; + try (ResultSet rs = statement.executeQuery(query)) { + assertTrue(rs.next()); + assertEquals(1, rs.getInt("val")); + } + + Dataset dataset = bigQuery.getDataset(DatasetId.of(tempDatasetName)); + assertNotNull(dataset); + assertEquals(location, dataset.getLocation()); + } finally { + bigQuery.delete(DatasetId.of(tempDatasetName), BigQuery.DatasetDeleteOption.deleteContents()); + } + } } From a3dfacd7468dd46b4a1fe9617640370d1dbb0c07 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Wed, 17 Jun 2026 13:09:11 -0400 Subject: [PATCH 2/2] nit --- .../com/google/cloud/bigquery/jdbc/it/ITStatementTest.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java index 2aaa750b3ea0..5aa5d6e03f4a 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java @@ -446,7 +446,12 @@ public void testTemporaryDatasetLocation() throws SQLException, InterruptedExcep assertNotNull(dataset); assertEquals(location, dataset.getLocation()); } finally { - bigQuery.delete(DatasetId.of(tempDatasetName), BigQuery.DatasetDeleteOption.deleteContents()); + try { + bigQuery.delete( + DatasetId.of(tempDatasetName), BigQuery.DatasetDeleteOption.deleteContents()); + } catch (Exception e) { + // Ignore cleanup exceptions to avoid masking the primary test failure + } } } }