From 98aa01306aac0d0bc853da015e392d81f47b58db Mon Sep 17 00:00:00 2001 From: Kirill Logachev Date: Wed, 6 May 2026 04:49:34 +0000 Subject: [PATCH 1/2] fix(bqjdbc): fallback to RestAPI if ReadAPI is not accessible --- .../bigquery/jdbc/BigQueryStatement.java | 61 +++++++++----- .../bigquery/jdbc/BigQueryStatementTest.java | 79 +++++++++++++++++++ 2 files changed, 119 insertions(+), 21 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 2c04747f8863..2bbf07b2d548 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -18,6 +18,8 @@ import com.google.api.core.InternalApi; import com.google.api.gax.paging.Page; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.StatusCode; import com.google.cloud.Tuple; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.JobListOption; @@ -57,6 +59,8 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterators; import com.google.common.util.concurrent.Uninterruptibles; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; import java.lang.ref.ReferenceQueue; import java.sql.Connection; import java.sql.ResultSet; @@ -879,9 +883,8 @@ Thread populateArrowBufferedQueue( rowsRead += response.getRowCount(); } break; - } catch (com.google.api.gax.rpc.ApiException e) { - if (e.getStatusCode().getCode() - == com.google.api.gax.rpc.StatusCode.Code.NOT_FOUND) { + } catch (ApiException e) { + if (e.getStatusCode().getCode() == StatusCode.Code.NOT_FOUND) { LOG.warning("Read session expired or not found: %s", e.getMessage()); enqueueError(arrowBatchWrapperBlockingQueue, e); break; @@ -929,25 +932,44 @@ Thread populateArrowBufferedQueue( /** Executes SQL query using either fast query path or read API */ void processQueryResponse(String query, TableResult results) throws SQLException { - LOG.finest( - "API call completed{Query=%s, Parent Job ID=%s, Total rows=%s} ", - query, results.getJobId(), results.getTotalRows()); - JobId currentJobId = results.getJobId(); - if (currentJobId == null) { - LOG.fine("Standard API with Stateless query used."); - this.currentResultSet = processJsonResultSet(results); - } else if (useReadAPI(results)) { - LOG.fine("HighThroughputAPI used."); - LOG.info("HTAPI job ID: " + currentJobId.getJob()); - this.currentResultSet = processArrowResultSet(results); - } else { - // read API cannot be used. - LOG.fine("Standard API used."); - this.currentResultSet = processJsonResultSet(results); + JobId jobId = results.getJobId(); + String queryId = results.getQueryId(); + LOG.info( + "Processing query response. JobId: %s, QueryId: %s, Total rows: %s", + jobId, queryId, results.getTotalRows()); + LOG.fine("Processing query response. Query: %s} ", query); + + ResultSet resultSet = null; + if (jobId != null && useReadAPI(results)) { + try { + LOG.info("Using ReadAPI to read the data."); + resultSet = processArrowResultSet(results); + } catch (SQLException e) { + if (!isPermissionDeniedException(e)) { + throw e; + } + LOG.log(Level.WARNING, "Permission denied for Read API, falling back to JSON API", e); + } } + + if (resultSet == null) { + LOG.info("Using Standard API to read the data."); + resultSet = processJsonResultSet(results); + } + this.currentResultSet = resultSet; this.currentUpdateCount = -1; } + private boolean isPermissionDeniedException(Throwable t) { + if (t == null) { + return false; + } + if (t instanceof StatusRuntimeException) { + return ((StatusRuntimeException) t).getStatus().getCode() == Status.Code.PERMISSION_DENIED; + } + return isPermissionDeniedException(t.getCause()); + } + // The read Ratio should be met // AND the User must not have disabled the Read API @VisibleForTesting @@ -977,9 +999,6 @@ private boolean meetsReadRatio(TableResult results) { } BigQueryJsonResultSet processJsonResultSet(TableResult results) { - String jobIdOrQueryId = - results.getJobId() == null ? results.getQueryId() : results.getJobId().getJob(); - LOG.info("BigQuery Job %s completed. Fetching results.", jobIdOrQueryId); List threadList = new ArrayList(); Schema schema = results.getSchema(); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 9fef90c69a4d..a7469f3cff26 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -26,6 +26,8 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.StatusCode; import com.google.cloud.ServiceOptions; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.QueryResultsOption; @@ -44,6 +46,7 @@ import com.google.cloud.bigquery.StandardSQLTypeName; import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableResult; +import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.jdbc.BigQueryStatement.JobIdWrapper; import com.google.cloud.bigquery.spi.BigQueryRpcFactory; import com.google.cloud.bigquery.storage.v1.ArrowSchema; @@ -65,6 +68,7 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; @@ -494,4 +498,79 @@ public void testGetStatementType(boolean isReadOnlyTokenUsed) throws Exception { verify(bigquery, isReadOnlyTokenUsed ? Mockito.never() : Mockito.times(1)) .create(any(JobInfo.class)); } + + @Test + public void testProcessQueryResponseFallbackToJsonOnReadApiFailure() throws SQLException { + BigQueryStatement statementSpy = Mockito.spy(bigQueryStatement); + TableResult tableResultMock = mockTableResultWithJob("job-id"); + + // Force useReadAPI to return true to enter the HTAPI block + doReturn(true).when(statementSpy).useReadAPI(tableResultMock); + + // Mock a permission denied ApiException + ApiException apiExceptionMock = mockApiException(StatusCode.Code.PERMISSION_DENIED); + + BigQueryJdbcException exceptionToThrow = + new BigQueryJdbcException("Simulated permission denied", apiExceptionMock); + + // Force processArrowResultSet to throw the permission exception + Mockito.doThrow(exceptionToThrow).when(statementSpy).processArrowResultSet(tableResultMock); + + BigQueryJsonResultSet jsonResultSetMock = mock(BigQueryJsonResultSet.class); + // Mock processJsonResultSet to return our mock JSON result set + doReturn(jsonResultSetMock).when(statementSpy).processJsonResultSet(tableResultMock); + + statementSpy.processQueryResponse("SELECT 1", tableResultMock); + + // Verify that processJsonResultSet was indeed called as a fallback + verify(statementSpy).processJsonResultSet(tableResultMock); + // Verify that currentResultSet is set to the mocked JSON result set + assertThat(statementSpy.currentResultSet).isEqualTo(jsonResultSetMock); + } + + @Test + public void testProcessQueryResponseNoFallbackOnNonPermissionFailure() throws SQLException { + BigQueryStatement statementSpy = Mockito.spy(bigQueryStatement); + TableResult tableResultMock = mockTableResultWithJob("job-id"); + + // Force useReadAPI to return true to enter the HTAPI block + doReturn(true).when(statementSpy).useReadAPI(tableResultMock); + + // Mock a non-permission ApiException (e.g., INTERNAL) + ApiException apiExceptionMock = mockApiException(StatusCode.Code.INTERNAL); + + BigQueryJdbcException exceptionToThrow = + new BigQueryJdbcException("Simulated internal error", apiExceptionMock); + + // Force processArrowResultSet to throw the non-permission exception + Mockito.doThrow(exceptionToThrow).when(statementSpy).processArrowResultSet(tableResultMock); + + BigQueryJsonResultSet jsonResultSetMock = mock(BigQueryJsonResultSet.class); + doReturn(jsonResultSetMock).when(statementSpy).processJsonResultSet(tableResultMock); + + // Assert that the exception is propagated + try { + statementSpy.processQueryResponse("SELECT 1", tableResultMock); + Assertions.fail("Expected SQLException to be thrown"); + } catch (SQLException e) { + assertEquals(exceptionToThrow, e); + } + + // Verify that processJsonResultSet was NOT called + verify(statementSpy, Mockito.never()).processJsonResultSet(tableResultMock); + } + + private TableResult mockTableResultWithJob(String jobId) { + TableResult tableResult = mock(TableResult.class); + doReturn(JobId.of(jobId)).when(tableResult).getJobId(); + return tableResult; + } + + private ApiException mockApiException(StatusCode.Code code) { + ApiException apiExceptionMock = mock(ApiException.class); + StatusCode statusCodeMock = mock(StatusCode.class); + doReturn(statusCodeMock).when(apiExceptionMock).getStatusCode(); + doReturn(code).when(statusCodeMock).getCode(); + return apiExceptionMock; + } } From 0989414f59c31bc52d4f00e25682338a0808261b Mon Sep 17 00:00:00 2001 From: Kirill Logachev Date: Wed, 6 May 2026 05:02:18 +0000 Subject: [PATCH 2/2] feedback --- .../cloud/bigquery/jdbc/BigQueryStatement.java | 15 +++++++++------ .../bigquery/jdbc/BigQueryStatementTest.java | 4 ++-- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 2bbf07b2d548..51925a4ef8a3 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -961,13 +961,16 @@ void processQueryResponse(String query, TableResult results) throws SQLException } private boolean isPermissionDeniedException(Throwable t) { - if (t == null) { - return false; - } - if (t instanceof StatusRuntimeException) { - return ((StatusRuntimeException) t).getStatus().getCode() == Status.Code.PERMISSION_DENIED; + while (t != null) { + if (t instanceof StatusRuntimeException) { + return ((StatusRuntimeException) t).getStatus().getCode() == Status.Code.PERMISSION_DENIED; + } + if (t instanceof ApiException) { + return ((ApiException) t).getStatusCode().getCode() == StatusCode.Code.PERMISSION_DENIED; + } + t = t.getCause(); } - return isPermissionDeniedException(t.getCause()); + return false; } // The read Ratio should be met diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index a7469f3cff26..af68e1198831 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -20,6 +20,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doReturn; @@ -68,7 +69,6 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; @@ -551,7 +551,7 @@ public void testProcessQueryResponseNoFallbackOnNonPermissionFailure() throws SQ // Assert that the exception is propagated try { statementSpy.processQueryResponse("SELECT 1", tableResultMock); - Assertions.fail("Expected SQLException to be thrown"); + fail("Expected SQLException to be thrown"); } catch (SQLException e) { assertEquals(exceptionToThrow, e); }