diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 239354ce1c92..2c151a74e14e 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -41,9 +41,9 @@ import com.google.common.base.Function; import com.google.common.base.Strings; import com.google.common.base.Supplier; +import com.google.common.collect.FluentIterable; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; -import com.google.common.collect.Lists; import com.google.common.collect.Maps; import java.util.List; import java.util.Map; @@ -548,17 +548,18 @@ public InsertAllResponse insertAll(InsertAllRequest request) { // an anonymous inner class. final boolean[] allInsertIdsSet = {true}; List rowsPb = - Lists.transform( - request.getRows(), - new Function() { - @Override - public Rows apply(RowToInsert rowToInsert) { - allInsertIdsSet[0] &= rowToInsert.getId() != null; - return new Rows() - .setInsertId(rowToInsert.getId()) - .setJson(rowToInsert.getContent()); - } - }); + FluentIterable.from(request.getRows()) + .transform( + new Function() { + @Override + public Rows apply(RowToInsert rowToInsert) { + allInsertIdsSet[0] &= rowToInsert.getId() != null; + return new Rows() + .setInsertId(rowToInsert.getId()) + .setJson(rowToInsert.getContent()); + } + }) + .toList(); requestPb.setRows(rowsPb); TableDataInsertAllResponse responsePb; diff --git a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 34498fa21a58..5d274bead3ec 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -801,7 +801,7 @@ public void testUpdateTableWithSelectedFields() { } @Test - public void testInsertAll() { + public void testInsertAllWithRowIdShouldRetry() { Map row1 = ImmutableMap.of("field", "value1"); Map row2 = ImmutableMap.of("field", "value2"); List rows = @@ -836,10 +836,17 @@ public TableDataInsertAllRequest.Rows apply(RowToInsert rowToInsert) { new TableDataInsertAllResponse.InsertErrors() .setIndex(0L) .setErrors(ImmutableList.of(new ErrorProto().setMessage("ErrorMessage"))))); + EasyMock.expect(bigqueryRpcMock.insertAll(PROJECT, DATASET, TABLE, requestPb)) + .andThrow(new BigQueryException(500, "InternalError")); EasyMock.expect(bigqueryRpcMock.insertAll(PROJECT, DATASET, TABLE, requestPb)) .andReturn(responsePb); EasyMock.replay(bigqueryRpcMock); - bigquery = options.getService(); + bigquery = + options + .toBuilder() + .setRetrySettings(ServiceOptions.getDefaultRetrySettings()) + .build() + .getService(); InsertAllResponse response = bigquery.insertAll(request); assertNotNull(response.getErrorsFor(0L)); assertNull(response.getErrorsFor(1L)); @@ -847,6 +854,47 @@ public TableDataInsertAllRequest.Rows apply(RowToInsert rowToInsert) { assertEquals("ErrorMessage", response.getErrorsFor(0L).get(0).getMessage()); } + @Test + public void testInsertAllWithoutRowIdShouldNotRetry() { + Map row1 = ImmutableMap.of("field", "value1"); + Map row2 = ImmutableMap.of("field", "value2"); + List rows = ImmutableList.of(RowToInsert.of(row1), RowToInsert.of(row2)); + InsertAllRequest request = + InsertAllRequest.newBuilder(TABLE_ID) + .setRows(rows) + .setSkipInvalidRows(false) + .setIgnoreUnknownValues(true) + .setTemplateSuffix("suffix") + .build(); + TableDataInsertAllRequest requestPb = + new TableDataInsertAllRequest() + .setRows( + Lists.transform( + rows, + new Function() { + @Override + public TableDataInsertAllRequest.Rows apply(RowToInsert rowToInsert) { + return new TableDataInsertAllRequest.Rows() + .setInsertId(rowToInsert.getId()) + .setJson(rowToInsert.getContent()); + } + })) + .setSkipInvalidRows(false) + .setIgnoreUnknownValues(true) + .setTemplateSuffix("suffix"); + EasyMock.expect(bigqueryRpcMock.insertAll(PROJECT, DATASET, TABLE, requestPb)) + .andThrow(new BigQueryException(500, "InternalError")); + EasyMock.replay(bigqueryRpcMock); + bigquery = + options + .toBuilder() + .setRetrySettings(ServiceOptions.getDefaultRetrySettings()) + .build() + .getService(); + thrown.expect(BigQueryException.class); + bigquery.insertAll(request); + } + @Test public void testInsertAllWithProject() { Map row1 = ImmutableMap.of("field", "value1");