fix(bigquery-jdbc): explicitly set location when available for temporary dataset (#13508)

b/524850173

Author: Neenu1995

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -1419,11 +1419,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());
     }
   }
 

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;
@@ -711,4 +713,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());
+  }
 }

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<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
+      }
+    }
+  }
 }