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..cec7250eabfc 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,14 @@ 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()); + String location = this.connection.getLocation(); + if (location != null && !location.isEmpty()) { + datasetInfoBuilder.setLocation(location); + } + 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..eaafdd3ef0a3 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 @@ -24,7 +24,15 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import com.google.api.gax.paging.Page; 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 com.google.cloud.bigquery.FieldValueList; +import com.google.cloud.bigquery.Table; +import com.google.cloud.bigquery.TableResult; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; @@ -412,4 +420,59 @@ 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;LargeResultDataset=" + + 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()); + + // Validate that the query results were written to a table in this dataset + Page tables = dataset.list(); + boolean tableFound = false; + for (Table table : tables.iterateAll()) { + if (table.getTableId().getTable().startsWith("temp_table_")) { + tableFound = true; + TableResult tableData = bigQuery.listTableData(table.getTableId()); + assertNotNull(tableData); + assertEquals(1, tableData.getTotalRows()); + for (FieldValueList row : tableData.iterateAll()) { + assertEquals(1, row.get(0).getLongValue()); + } + break; + } + } + assertTrue(tableFound, "Expected temporary table was not found in the dataset"); + } finally { + try { + bigQuery.delete( + DatasetId.of(tempDatasetName), BigQuery.DatasetDeleteOption.deleteContents()); + } catch (Exception e) { + // Ignore cleanup exceptions to avoid masking the primary test failure + } + } + } }