diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java index 66d35f6c460e..62588c58ab42 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java @@ -29,13 +29,13 @@ class BigQueryArrowBatchWrapper { // Marks the end of the stream for the ResultSet private final boolean isLast; - private final Exception exception; + private final Throwable exception; private BigQueryArrowBatchWrapper( ArrowRecordBatch currentArrowBatch, JsonStringArrayList nestedRecords, boolean isLast, - Exception exception) { + Throwable exception) { this.currentArrowBatch = currentArrowBatch; this.nestedRecords = nestedRecords; this.isLast = isLast; @@ -53,7 +53,7 @@ static BigQueryArrowBatchWrapper getNestedFieldValueListWrapper( return new BigQueryArrowBatchWrapper(null, nestedRecords, isLastFlag, null); } - static BigQueryArrowBatchWrapper ofError(Exception exception) { + static BigQueryArrowBatchWrapper ofError(Throwable exception) { return new BigQueryArrowBatchWrapper(null, null, true, exception); } @@ -69,7 +69,7 @@ boolean isLast() { return this.isLast; } - Exception getException() { + Throwable getException() { return this.exception; } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java index a0d51a64503a..0317a24a0b86 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java @@ -854,7 +854,7 @@ public void clearWarnings() throws SQLException { this.warningsLoaded = true; } - static BigQueryJdbcException backgroundFetchError(Exception cause) { + static BigQueryJdbcException backgroundFetchError(Throwable cause) { String msg = "Error occurred while fetching query results"; return cause instanceof BigQueryException ? new BigQueryJdbcException(msg, (BigQueryException) cause) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryFieldValueListWrapper.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryFieldValueListWrapper.java index 840b9dcd347b..64d8ba75bfd7 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryFieldValueListWrapper.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryFieldValueListWrapper.java @@ -48,7 +48,7 @@ class BigQueryFieldValueListWrapper { // This flag marks the end of the stream for the ResultSet private boolean isLast = false; - private final Exception exception; + private final Throwable exception; static BigQueryFieldValueListWrapper ofEndOfStream(FieldList fieldList) { return new BigQueryFieldValueListWrapper(fieldList, null, null, null, true, null); @@ -106,7 +106,7 @@ static BigQueryFieldValueListWrapper getNestedFieldValueListWrapper( fieldList, null, arrayFieldValueList, null, isLastFlag, null); } - static BigQueryFieldValueListWrapper ofError(Exception exception) { + static BigQueryFieldValueListWrapper ofError(Throwable exception) { return new BigQueryFieldValueListWrapper(null, null, null, null, true, exception); } @@ -116,7 +116,7 @@ private BigQueryFieldValueListWrapper( List arrayFieldValueList, Object[] rowValues, boolean isLast, - Exception exception) { + Throwable exception) { this.fieldList = fieldList; this.fieldValueList = fieldValueList; this.arrayFieldValueList = arrayFieldValueList; @@ -165,7 +165,7 @@ boolean isLast() { return this.isLast; } - Exception getException() { + Throwable getException() { return this.exception; } } 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 b070356cb221..e036030e6d6f 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 @@ -1077,7 +1077,7 @@ Future populateArrowBufferedQueue( + " Interrupted @ arrowStreamProcessor", e); enqueueError(arrowBatchWrapperBlockingQueue, e); - } catch (Exception e) { + } catch (Throwable e) { if (e.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { Thread.currentThread().interrupt(); @@ -1090,11 +1090,14 @@ Future populateArrowBufferedQueue( e); enqueueError(arrowBatchWrapperBlockingQueue, e); } else { + enqueueError(arrowBatchWrapperBlockingQueue, e); LOG.log( Level.WARNING, "\n" + Thread.currentThread().getName() + " Error @ arrowStreamProcessor", e); - enqueueError(arrowBatchWrapperBlockingQueue, e); + } + if (e instanceof Error) { + throw (Error) e; } } finally { // logic needed for graceful shutdown enqueueEndOfStream(arrowBatchWrapperBlockingQueue); @@ -1322,11 +1325,11 @@ Future runNextPageTaskAsync( String currentPageToken = firstPageToken; TableResult currentResults = result; TableId destinationTable = null; - if (firstPageToken != null) { - destinationTable = getDestinationTable(jobId); - } try { + if (firstPageToken != null) { + destinationTable = getDestinationTable(jobId); + } while (currentPageToken != null) { // do not process further pages and shutdown if (isWorkerShuttingDown()) { @@ -1368,12 +1371,15 @@ Future runNextPageTaskAsync( "%s Interrupted @ runNextPageTaskAsync: %s", Thread.currentThread().getName(), ex.getMessage()); enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); - } catch (Exception ex) { + } catch (Throwable ex) { if (ex.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { Thread.currentThread().interrupt(); } enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); + if (ex instanceof Error) { + throw (Error) ex; + } } finally { // this will stop the parseDataTask as well when the pagination // completes @@ -1457,7 +1463,7 @@ Future parseAndPopulateRpcDataAsync( + " Interrupted @ populateBufferAsync", ex); enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); - } catch (Exception ex) { + } catch (Throwable ex) { if (ex.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) { Thread.currentThread().interrupt(); @@ -1470,11 +1476,11 @@ Future parseAndPopulateRpcDataAsync( ex); enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } else { + enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); LOG.log( Level.WARNING, "\n" + Thread.currentThread().getName() + " Error @ populateBufferAsync", ex); - enqueueBufferError(bigQueryFieldValueListWrapperBlockingQueue, ex); } } finally { enqueueBufferEndOfStream(bigQueryFieldValueListWrapperBlockingQueue); @@ -1979,7 +1985,7 @@ private void enqueueMarker(BlockingQueue queue, T marker) { } } - private void enqueueError(BlockingQueue queue, Exception e) { + private void enqueueError(BlockingQueue queue, Throwable e) { enqueueMarker(queue, BigQueryArrowBatchWrapper.ofError(e)); } @@ -1987,7 +1993,7 @@ private void enqueueEndOfStream(BlockingQueue queue) enqueueMarker(queue, BigQueryArrowBatchWrapper.of(null, true)); } - private void enqueueBufferError(BlockingQueue queue, Exception e) { + private void enqueueBufferError(BlockingQueue queue, Throwable e) { enqueueMarker(queue, BigQueryFieldValueListWrapper.ofError(e)); } 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 41f3731c2806..4ad8358f3894 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 @@ -47,6 +47,7 @@ import com.google.cloud.bigquery.DatasetInfo; import com.google.cloud.bigquery.Field; import com.google.cloud.bigquery.FieldList; +import com.google.cloud.bigquery.FieldValue; import com.google.cloud.bigquery.FieldValueList; import com.google.cloud.bigquery.Job; import com.google.cloud.bigquery.JobId; @@ -87,6 +88,7 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.Iterator; @@ -1560,4 +1562,157 @@ public ReadRowsResponse next() { marker.getException(), "Stopping early due to worker shutdown must enqueue an error marker rather than normal end-of-stream"); } + + private static FieldValueList createTestRow(boolean first, long second) { + return FieldValueList.of( + Arrays.asList( + FieldValue.of(FieldValue.Attribute.PRIMITIVE, Boolean.toString(first)), + FieldValue.of(FieldValue.Attribute.PRIMITIVE, Long.toString(second))), + fieldList); + } + + @Test + public void testRunNextPageTaskAsyncPropagatesOutOfMemoryErrorInsteadOfTruncating() + throws Exception { + JobId testJobId = JobId.of("project", "job_oom_next_page"); + TableId destTable = TableId.of("project", "dataset", "dest_table"); + + Job mockJob = mock(Job.class); + QueryJobConfiguration jobConfig = + QueryJobConfiguration.newBuilder("SELECT first, second FROM test") + .setDestinationTable(destTable) + .build(); + doReturn(mockJob).when(bigquery).getJob(testJobId); + doReturn(jobConfig).when(mockJob).getConfiguration(); + + Schema schema = Schema.of(fieldList); + TableResult firstPage = mock(TableResult.class); + doReturn(schema).when(firstPage).getSchema(); + doReturn(10L).when(firstPage).getTotalRows(); + doReturn(testJobId).when(firstPage).getJobId(); + doReturn("page_2_token").when(firstPage).getNextPageToken(); + doReturn(Arrays.asList(createTestRow(true, 1L), createTestRow(false, 2L))) + .when(firstPage) + .getValues(); + + Mockito.doThrow(new OutOfMemoryError("Simulated OOM allocating TableResult for page 2")) + .when(bigquery) + .listTableData( + any(TableId.class), any(TableDataListOption.class), any(TableDataListOption.class)); + + try (ResultSet rs = bigQueryStatement.processJsonResultSet(firstPage, mockJob)) { + SQLException thrown = + assertThrows( + SQLException.class, + () -> { + while (rs.next()) { + // consume until error + } + }); + assertTrue(thrown.getCause() instanceof OutOfMemoryError); + } + } + + @Test + public void testParseAndPopulateRpcDataAsyncPropagatesOutOfMemoryErrorInsteadOfTruncating() + throws Exception { + Iterable oomOnThirdRowIterable = + () -> + new Iterator() { + private int index = 0; + + @Override + public boolean hasNext() { + return index < 10; + } + + @Override + public FieldValueList next() { + index++; + if (index <= 2) { + return createTestRow(true, index); + } + throw new OutOfMemoryError("Simulated OOM unpacking row " + index); + } + }; + + Schema schema = Schema.of(fieldList); + TableResult firstPage = mock(TableResult.class); + doReturn(schema).when(firstPage).getSchema(); + doReturn(10L).when(firstPage).getTotalRows(); + doReturn(null).when(firstPage).getJobId(); + doReturn(oomOnThirdRowIterable).when(firstPage).getValues(); + + try (ResultSet rs = bigQueryStatement.processJsonResultSet(firstPage, null)) { + SQLException thrown = + assertThrows( + SQLException.class, + () -> { + while (rs.next()) { + // consume until error + } + }); + assertTrue(thrown.getCause() instanceof OutOfMemoryError); + } + } + + @Test + public void testPopulateArrowBufferedQueuePropagatesOutOfMemoryErrorInsteadOfTruncating() + throws Exception { + BlockingQueue boundedQueue = new LinkedBlockingDeque<>(4); + ReadSession readSession = + ReadSession.newBuilder() + .setArrowSchema(arrowSchema) + .addStreams( + ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/0")) + .build(); + + @SuppressWarnings("unchecked") + ServerStream serverStream = mock(ServerStream.class); + Iterator responseIterator = + new Iterator() { + @Override + public boolean hasNext() { + return true; + } + + @Override + public ReadRowsResponse next() { + throw new OutOfMemoryError("Simulated OOM reading Arrow batch"); + } + }; + doReturn(responseIterator).when(serverStream).iterator(); + + @SuppressWarnings("unchecked") + ServerStreamingCallable callable = + mock(ServerStreamingCallable.class); + doReturn(serverStream).when(callable).call(any(ReadRowsRequest.class)); + + BigQueryReadClient readClient = mock(BigQueryReadClient.class); + doReturn(callable).when(readClient).readRowsCallable(); + + Future producerFuture = + bigQueryStatement.populateArrowBufferedQueue(readSession, boundedQueue, readClient); + + try (BigQueryArrowResultSet rs = + BigQueryArrowResultSet.of( + Schema.of(fieldList), + arrowSchema, + 10L, + bigQueryStatement, + boundedQueue, + producerFuture, + bigquery)) { + SQLException thrown = + assertThrows( + SQLException.class, + () -> { + while (rs.next()) { + // consume until error + } + }); + assertNotNull(thrown.getCause()); + assertTrue(thrown.getCause().getCause() instanceof OutOfMemoryError); + } + } }