From 869d5d4c0cec02947c1e115a5e8c9f5a57de69c7 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Thu, 20 Aug 2026 21:13:33 +0000 Subject: [PATCH 1/3] chore(bigquery-jdbc): eliminate dry run to resolve statement type --- .../bigquery/jdbc/BigQueryStatement.java | 30 +++++++---- .../bigquery/jdbc/BigQueryStatementTest.java | 53 ++++++++++++++----- 2 files changed, 60 insertions(+), 23 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 30c11d13d263..481417c1837f 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 @@ -602,6 +602,16 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration) return new ExecuteResult(tableResult, job); } + private StatementType getStatementType(ExecuteResult executeResult) { + if (executeResult.tableResult != null && executeResult.tableResult.getStatementType() != null) { + return executeResult.tableResult.getStatementType(); + } + if (executeResult.job != null && executeResult.job.getStatistics() instanceof QueryStatistics) { + return ((QueryStatistics) executeResult.job.getStatistics()).getStatementType(); + } + return null; + } + /** * Execute the SQL script and sets the reference of the underlying job, passing null querySettings * will result in the FastQueryPath @@ -620,10 +630,7 @@ void runQuery(String query, QueryJobConfiguration jobConfiguration) try { resetStatementFields(); ExecuteResult executeResult = executeJob(jobConfiguration); - StatementType statementType = - executeResult.job == null - ? getStatementType(jobConfiguration) - : ((QueryStatistics) executeResult.job.getStatistics()).getStatementType(); + StatementType statementType = getStatementType(executeResult); SqlType queryType = getQueryType(jobConfiguration, statementType); handleQueryResult(query, executeResult.tableResult, queryType, executeResult.job); } catch (InterruptedException ex) { @@ -709,11 +716,16 @@ void handleQueryResult(String query, TableResult results, SqlType queryType, Job break; case DML: case DML_EXTRA: - QueryStatistics dmlStats = getQueryStatisticsFromJob(results, job); - Long dmlRowCount = - (dmlStats != null && dmlStats.getNumDmlAffectedRows() != null) - ? dmlStats.getNumDmlAffectedRows() - : 0L; + Long dmlRowCount; + if (results != null && results.getNumDmlAffectedRows() != null) { + dmlRowCount = results.getNumDmlAffectedRows(); + } else { + QueryStatistics dmlStats = getQueryStatisticsFromJob(results, job); + dmlRowCount = + (dmlStats != null && dmlStats.getNumDmlAffectedRows() != null) + ? dmlStats.getNumDmlAffectedRows() + : 0L; + } updateAffectedRowCount(dmlRowCount); break; case TCL: 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 91ae858a6ab8..23d686cbab67 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 @@ -20,6 +20,8 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; @@ -169,6 +171,10 @@ private TableResult setupMockQueryResults(JobId jobId, StatementType type, Long TableResult tableResultMock = mock(TableResult.class); doReturn(jobId).when(tableResultMock).getJobId(); doReturn(Schema.of()).when(tableResultMock).getSchema(); + doReturn(type).when(tableResultMock).getStatementType(); + if (affectedRows != null) { + doReturn(affectedRows).when(tableResultMock).getNumDmlAffectedRows(); + } doReturn(tableResultMock) .when(bigquery) .queryWithTimeout(any(QueryJobConfiguration.class), any(), any()); @@ -447,6 +453,7 @@ public void testJoblessQuery() throws SQLException, InterruptedException { TableResult tableResultMock = mock(TableResult.class); doReturn("queryId").when(tableResultMock).getQueryId(); doReturn(null).when(tableResultMock).getJobId(); + doReturn(StatementType.SELECT).when(tableResultMock).getStatementType(); doReturn(tableResultMock) .when(bigquery) .queryWithTimeout(any(QueryJobConfiguration.class), any(), any()); @@ -454,17 +461,10 @@ public void testJoblessQuery() throws SQLException, InterruptedException { .when(joblessStatementSpy) .processJsonResultSet(eq(tableResultMock), any()); - Job dryRunJobMock = getJobMock(null, null, StatementType.SELECT); - ArgumentCaptor dryRunCaptor = ArgumentCaptor.forClass(JobInfo.class); - doReturn(dryRunJobMock).when(bigquery).create(dryRunCaptor.capture()); - joblessStatementSpy.executeQuery("SELECT 1"); verify(bigquery).queryWithTimeout(any(QueryJobConfiguration.class), any(), any()); - verify(bigquery).create(any(JobInfo.class)); - assertTrue( - Boolean.TRUE.equals( - ((QueryJobConfiguration) dryRunCaptor.getValue().getConfiguration()).dryRun())); + verify(bigquery, Mockito.never()).create(any(JobInfo.class)); // 2. Test JobCreationMode=1 (jobful) Mockito.reset(bigquery); @@ -914,6 +914,7 @@ public void testExecute_propagatesContextAndBaggage() throws Exception { TableResult tableResultMock = mock(TableResult.class); doReturn(jobId).when(tableResultMock).getJobId(); doReturn(Schema.of()).when(tableResultMock).getSchema(); + doReturn(StatementType.SELECT).when(tableResultMock).getStatementType(); return tableResultMock; }) .when(bigquery) @@ -923,8 +924,6 @@ public void testExecute_propagatesContextAndBaggage() throws Exception { // Setup connection mocks to allow the statement to execute successfully doReturn(true).when(bigQueryConnection).getUseStatelessQueryMode(); - Job dryRunJobMock = getJobMock(null, null, StatementType.SELECT); - doReturn(dryRunJobMock).when(bigquery).create(Mockito.any(JobInfo.class)); BigQueryJsonResultSet resultSetMock = mock(BigQueryJsonResultSet.class); doReturn(resultSetMock) @@ -937,6 +936,7 @@ public void testExecute_propagatesContextAndBaggage() throws Exception { // Verify the SDK call actually occurred verify(bigquery) .queryWithTimeout(Mockito.any(QueryJobConfiguration.class), Mockito.any(), Mockito.any()); + verify(bigquery, Mockito.never()).create(Mockito.any(JobInfo.class)); } @Test @@ -1065,13 +1065,10 @@ public void testTemporaryDatasetCreationRespectsConnectionLocation() // 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(StatementType.SELECT).when(result).getStatementType(); doReturn(result) .when(bigquery) .queryWithTimeout(any(QueryJobConfiguration.class), any(JobId.class), any()); @@ -1091,4 +1088,32 @@ public void testTemporaryDatasetCreationRespectsConnectionLocation() assertEquals("temp_dataset", createdDatasetInfo.getDatasetId().getDataset()); assertEquals("europe-west3", createdDatasetInfo.getLocation()); } + + @Test + public void testStatelessQueryExecutionDoesNotInvokeDryRun() throws Exception { + TableResult tableResultMock = setupMockQueryResults(null, StatementType.SELECT, null); + BigQueryStatement statementSpy = Mockito.spy(bigQueryStatement); + doReturn(mock(BigQueryJsonResultSet.class)) + .when(statementSpy) + .processJsonResultSet(eq(tableResultMock), any()); + + boolean hasResultSet = statementSpy.execute("SELECT 1"); + + assertTrue(hasResultSet); + assertNotNull(statementSpy.getResultSet()); + verify(bigquery, Mockito.never()).create(any(JobInfo.class)); + } + + @Test + public void testStatelessDmlExecutionUsesTableResultWithoutDryRunOrGetJob() throws Exception { + setupMockQueryResults(null, StatementType.UPDATE, 15L); + + int updatedCount = bigQueryStatement.executeUpdate("UPDATE dataset.table SET col = 1"); + + assertEquals(15, updatedCount); + assertEquals(15L, bigQueryStatement.getLargeUpdateCount()); + assertNull(bigQueryStatement.getResultSet()); + verify(bigquery, Mockito.never()).create(any(JobInfo.class)); + verify(bigquery, Mockito.never()).getJob(any(JobId.class)); + } } From 93c5552a513293b977445159fe725a05046e32b2 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Thu, 20 Aug 2026 22:32:25 +0000 Subject: [PATCH 2/3] refactor(bigquery-jdbc): resolve statement type and DML stats via JobStatistics on TableResult --- .../bigquery/jdbc/BigQueryStatement.java | 14 +++++++----- .../bigquery/jdbc/BigQueryStatementTest.java | 22 ++++++++++++++----- 2 files changed, 24 insertions(+), 12 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 481417c1837f..9b67ca86fa6e 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 @@ -603,8 +603,9 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration) } private StatementType getStatementType(ExecuteResult executeResult) { - if (executeResult.tableResult != null && executeResult.tableResult.getStatementType() != null) { - return executeResult.tableResult.getStatementType(); + if (executeResult.tableResult != null + && executeResult.tableResult.getJobStatistics() instanceof QueryStatistics) { + return ((QueryStatistics) executeResult.tableResult.getJobStatistics()).getStatementType(); } if (executeResult.job != null && executeResult.job.getStatistics() instanceof QueryStatistics) { return ((QueryStatistics) executeResult.job.getStatistics()).getStatementType(); @@ -716,10 +717,11 @@ void handleQueryResult(String query, TableResult results, SqlType queryType, Job break; case DML: case DML_EXTRA: - Long dmlRowCount; - if (results != null && results.getNumDmlAffectedRows() != null) { - dmlRowCount = results.getNumDmlAffectedRows(); - } else { + Long dmlRowCount = null; + if (results != null && results.getJobStatistics() instanceof QueryStatistics) { + dmlRowCount = ((QueryStatistics) results.getJobStatistics()).getNumDmlAffectedRows(); + } + if (dmlRowCount == null) { QueryStatistics dmlStats = getQueryStatisticsFromJob(results, job); dmlRowCount = (dmlStats != null && dmlStats.getNumDmlAffectedRows() != null) 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 23d686cbab67..f56136b04a56 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 @@ -171,9 +171,13 @@ private TableResult setupMockQueryResults(JobId jobId, StatementType type, Long TableResult tableResultMock = mock(TableResult.class); doReturn(jobId).when(tableResultMock).getJobId(); doReturn(Schema.of()).when(tableResultMock).getSchema(); - doReturn(type).when(tableResultMock).getStatementType(); - if (affectedRows != null) { - doReturn(affectedRows).when(tableResultMock).getNumDmlAffectedRows(); + if (type != null || affectedRows != null) { + QueryStatistics queryStatsMock = mock(QueryStatistics.class); + doReturn(type).when(queryStatsMock).getStatementType(); + if (affectedRows != null) { + doReturn(affectedRows).when(queryStatsMock).getNumDmlAffectedRows(); + } + doReturn(queryStatsMock).when(tableResultMock).getJobStatistics(); } doReturn(tableResultMock) .when(bigquery) @@ -453,7 +457,9 @@ public void testJoblessQuery() throws SQLException, InterruptedException { TableResult tableResultMock = mock(TableResult.class); doReturn("queryId").when(tableResultMock).getQueryId(); doReturn(null).when(tableResultMock).getJobId(); - doReturn(StatementType.SELECT).when(tableResultMock).getStatementType(); + QueryStatistics queryStatsMock = mock(QueryStatistics.class); + doReturn(StatementType.SELECT).when(queryStatsMock).getStatementType(); + doReturn(queryStatsMock).when(tableResultMock).getJobStatistics(); doReturn(tableResultMock) .when(bigquery) .queryWithTimeout(any(QueryJobConfiguration.class), any(), any()); @@ -914,7 +920,9 @@ public void testExecute_propagatesContextAndBaggage() throws Exception { TableResult tableResultMock = mock(TableResult.class); doReturn(jobId).when(tableResultMock).getJobId(); doReturn(Schema.of()).when(tableResultMock).getSchema(); - doReturn(StatementType.SELECT).when(tableResultMock).getStatementType(); + QueryStatistics queryStatsMock = mock(QueryStatistics.class); + doReturn(StatementType.SELECT).when(queryStatsMock).getStatementType(); + doReturn(queryStatsMock).when(tableResultMock).getJobStatistics(); return tableResultMock; }) .when(bigquery) @@ -1068,7 +1076,9 @@ public void testTemporaryDatasetCreationRespectsConnectionLocation() // 3. Mock bigquery.queryWithTimeout(...) to return tableResult (so execution doesn't fail on // query execution) TableResult result = mock(TableResult.class); - doReturn(StatementType.SELECT).when(result).getStatementType(); + QueryStatistics queryStatsMock = mock(QueryStatistics.class); + doReturn(StatementType.SELECT).when(queryStatsMock).getStatementType(); + doReturn(queryStatsMock).when(result).getJobStatistics(); doReturn(result) .when(bigquery) .queryWithTimeout(any(QueryJobConfiguration.class), any(JobId.class), any()); From 0e9ae278b05152f3a36adb4fb11563e8bee7cfb4 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Thu, 20 Aug 2026 22:40:33 +0000 Subject: [PATCH 3/3] address pr feedback --- .../google/cloud/bigquery/jdbc/BigQueryStatement.java | 9 +++++---- 1 file changed, 5 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 9b67ca86fa6e..878295f3a8af 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 @@ -717,11 +717,12 @@ void handleQueryResult(String query, TableResult results, SqlType queryType, Job break; case DML: case DML_EXTRA: - Long dmlRowCount = null; + Long dmlRowCount; if (results != null && results.getJobStatistics() instanceof QueryStatistics) { - dmlRowCount = ((QueryStatistics) results.getJobStatistics()).getNumDmlAffectedRows(); - } - if (dmlRowCount == null) { + Long affectedRows = + ((QueryStatistics) results.getJobStatistics()).getNumDmlAffectedRows(); + dmlRowCount = affectedRows != null ? affectedRows : 0L; + } else { QueryStatistics dmlStats = getQueryStatisticsFromJob(results, job); dmlRowCount = (dmlStats != null && dmlStats.getNumDmlAffectedRows() != null)