Merge branch 'main' into jdbc/fix-time-precision

Author: keshavdandeva

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

@@ -104,6 +104,7 @@ public class BigQueryStatement extends BigQueryNoOpsStatement {
 
   private final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString());
   public static final int DEFAULT_BUFFER_SIZE = BigQuerySettings.DEFAULT_NUM_BUFFERED_ROWS * 2;
+  static final int DEFAULT_ARROW_BUFFER_SIZE = 32;
   private static final String DEFAULT_DATASET_NAME = "_google_jdbc";
   private static final String DEFAULT_TABLE_NAME = "temp_table_";
   private static final String JDBC_JOB_PREFIX = "google-jdbc-";
@@ -931,7 +932,7 @@ ResultSet processArrowResultSet(TableResult results, Job job) throws SQLExceptio
               .setMaxStreamCount(1);
 
       ReadSession readSession = getReadSession(builder.build());
-      this.arrowBatchWrapperBlockingQueue = new LinkedBlockingDeque<>(getBufferSize());
+      this.arrowBatchWrapperBlockingQueue = new LinkedBlockingDeque<>(getArrowBufferSize());
       // deserialize and populate the buffer async, so that the client isn't blocked
       populateBufferWorker =
           populateArrowBufferedQueue(
@@ -1512,6 +1513,10 @@ private int getBufferSize() {
         : Math.min(this.querySettings.getNumBufferedRows() * 2, 100000));
   }
 
+  int getArrowBufferSize() {
+    return DEFAULT_ARROW_BUFFER_SIZE;
+  }
+
   /** Returns the destinationTable from jobId by calling `jobs.get` API */
   TableId getDestinationTable(JobId jobId) {
     Job job = this.bigQuery.getJob(jobId);

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java

@@ -1318,4 +1318,33 @@ public void testJoblessQueryPopulatesQueryStatistics() throws Exception {
 
     verify(bigquery, Mockito.never()).getJob(any(JobId.class));
   }
+
+  @Test
+  public void testProcessArrowResultSetUsesBoundedQueueCapacity() throws SQLException {
+    BigQueryStatement defaultStmtSpy = Mockito.spy(bigQueryStatement);
+    BigQueryReadClient bigQueryReadClient = mock(BigQueryReadClient.class);
+    Schema schema = Schema.of(fieldList);
+    ReadSession readSession = ReadSession.getDefaultInstance();
+    doReturn(bigQueryReadClient).when(defaultStmtSpy).getBigQueryReadClient();
+    doReturn(readSession).when(defaultStmtSpy).getReadSession(any(CreateReadSessionRequest.class));
+    Future<?> mockWorker = mock(Future.class);
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<BlockingQueue<BigQueryArrowBatchWrapper>> defaultQueueCaptor =
+        ArgumentCaptor.forClass(BlockingQueue.class);
+    doReturn(mockWorker)
+        .when(defaultStmtSpy)
+        .populateArrowBufferedQueue(
+            any(ReadSession.class), defaultQueueCaptor.capture(), any(BigQueryReadClient.class));
+    doReturn(arrowSchema).when(defaultStmtSpy).getArrowSchema(any(ReadSession.class));
+
+    TableResult result = mock(TableResult.class);
+    doReturn(schema).when(result).getSchema();
+    doReturn(10L).when(result).getTotalRows();
+    doReturn(TABLE_ID).when(defaultStmtSpy).getDestinationTable(any());
+
+    defaultStmtSpy.processArrowResultSet(result, null);
+    assertEquals(
+        BigQueryStatement.DEFAULT_ARROW_BUFFER_SIZE,
+        defaultQueueCaptor.getValue().remainingCapacity());
+  }
 }