Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<DatasetInfo> 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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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));
Comment thread
Neenu1995 marked this conversation as resolved.
assertNotNull(dataset);
assertEquals(location, dataset.getLocation());

// Validate that the query results were written to a table in this dataset
Page<Table> 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
}
}
Comment thread
Neenu1995 marked this conversation as resolved.
}
}
Loading