feat(jdbc): add retry mechanism for Arrow and JSON result streams

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

@@ -31,24 +31,35 @@ class BigQueryArrowBatchWrapper {
   // Marks the end of the stream for the ResultSet
   private final boolean isLast;
 
+  private final Exception exception;
+
   private BigQueryArrowBatchWrapper(
-      ArrowRecordBatch currentArrowBatch, JsonStringArrayList nestedRecords, boolean isLast) {
+      ArrowRecordBatch currentArrowBatch,
+      JsonStringArrayList nestedRecords,
+      boolean isLast,
+      Exception exception) {
     this.currentArrowBatch = currentArrowBatch;
     this.nestedRecords = nestedRecords;
     this.isLast = isLast;
+    this.exception = exception;
   }
 
   static BigQueryArrowBatchWrapper of(ArrowRecordBatch currentArrowBatch, boolean... isLast) {
     LOG.finest("++enter++");
     boolean isLastFlag = isLast != null && isLast.length == 1 && isLast[0];
-    return new BigQueryArrowBatchWrapper(currentArrowBatch, null, isLastFlag);
+    return new BigQueryArrowBatchWrapper(currentArrowBatch, null, isLastFlag, null);
   }
 
   static BigQueryArrowBatchWrapper getNestedFieldValueListWrapper(
       JsonStringArrayList nestedRecords, boolean... isLast) {
     LOG.finest("++enter++");
     boolean isLastFlag = isLast != null && isLast.length == 1 && isLast[0];
-    return new BigQueryArrowBatchWrapper(null, nestedRecords, isLastFlag);
+    return new BigQueryArrowBatchWrapper(null, nestedRecords, isLastFlag, null);
+  }
+
+  static BigQueryArrowBatchWrapper ofError(Exception exception) {
+    LOG.finest("++enter++");
+    return new BigQueryArrowBatchWrapper(null, null, true, exception);
   }
 
   ArrowRecordBatch getCurrentArrowBatch() {
@@ -65,4 +76,8 @@ boolean isLast() {
     LOG.finest("++enter++");
     return this.isLast;
   }
+
+  Exception getException() {
+    return this.exception;
+  }
 }

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

@@ -24,6 +24,7 @@
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.StandardSQLTypeName;
 import com.google.cloud.bigquery.exception.BigQueryJdbcException;
+import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
 import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
 import com.google.cloud.bigquery.storage.v1.ArrowSchema;
 import java.io.IOException;
@@ -236,6 +237,9 @@ public boolean next() throws SQLException {
           /* Start of iteration or we have exhausted the current batch */
           // Advance the cursor. Potentially blocking operation.
           BigQueryArrowBatchWrapper batchWrapper = this.buffer.take();
+          if (batchWrapper.getException() != null) {
+            throw new BigQueryJdbcRuntimeException(batchWrapper.getException());
+          }
           if (batchWrapper.isLast()) {
             /* Marks the end of the records */
             if (this.vectorSchemaRoot != null) {

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

@@ -39,28 +39,36 @@ class BigQueryFieldValueListWrapper {
 
   // This flag marks the end of the stream for the ResultSet
   private boolean isLast = false;
+  private final Exception exception;
 
   static BigQueryFieldValueListWrapper of(
       FieldList fieldList, FieldValueList fieldValueList, boolean... isLast) {
     boolean isLastFlag = isLast != null && isLast.length == 1 && isLast[0];
-    return new BigQueryFieldValueListWrapper(fieldList, fieldValueList, null, isLastFlag);
+    return new BigQueryFieldValueListWrapper(fieldList, fieldValueList, null, isLastFlag, null);
   }
 
   static BigQueryFieldValueListWrapper getNestedFieldValueListWrapper(
       FieldList fieldList, List<FieldValue> arrayFieldValueList, boolean... isLast) {
     boolean isLastFlag = isLast != null && isLast.length == 1 && isLast[0];
-    return new BigQueryFieldValueListWrapper(fieldList, null, arrayFieldValueList, isLastFlag);
+    return new BigQueryFieldValueListWrapper(
+        fieldList, null, arrayFieldValueList, isLastFlag, null);
+  }
+
+  static BigQueryFieldValueListWrapper ofError(Exception exception) {
+    return new BigQueryFieldValueListWrapper(null, null, null, true, exception);
   }
 
   private BigQueryFieldValueListWrapper(
       FieldList fieldList,
       FieldValueList fieldValueList,
       List<FieldValue> arrayFieldValueList,
-      boolean isLast) {
+      boolean isLast,
+      Exception exception) {
     this.fieldList = fieldList;
     this.fieldValueList = fieldValueList;
     this.arrayFieldValueList = arrayFieldValueList;
     this.isLast = isLast;
+    this.exception = exception;
   }
 
   public FieldList getFieldList() {
@@ -78,4 +86,8 @@ public List<FieldValue> getArrayFieldValueList() {
   public boolean isLast() {
     return this.isLast;
   }
+
+  public Exception getException() {
+    return this.exception;
+  }
 }

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

@@ -159,6 +159,9 @@ public boolean next() throws SQLException {
       try {
         // Advance the cursor,Potentially blocking operation
         this.cursor = this.buffer.take();
+        if (this.cursor.getException() != null) {
+          throw new BigQueryJdbcRuntimeException(this.cursor.getException());
+        }
         this.rowCnt++;
         // Check for end of stream
         if (this.cursor.isLast()) {