fix(bqjdbc): fallback to RestAPI if ReadAPI is not accessible

Author: logachev

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<Thread> threadList = new ArrayList<Thread>();
 
     Schema schema = results.getSchema();

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