Add retries for insertAll (#2518)

Author: vam-google

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -24,12 +24,14 @@
 import com.google.api.services.bigquery.model.GetQueryResultsResponse;
 import com.google.api.services.bigquery.model.TableDataInsertAllRequest;
 import com.google.api.services.bigquery.model.TableDataInsertAllRequest.Rows;
+import com.google.api.services.bigquery.model.TableDataInsertAllResponse;
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.api.services.bigquery.model.TableSchema;
 import com.google.cloud.BaseService;
 import com.google.cloud.PageImpl;
 import com.google.cloud.PageImpl.NextPageFetcher;
 import com.google.cloud.RetryHelper;
+import com.google.cloud.RetryHelper.RetryHelperException;
 import com.google.cloud.RetryOption;
 import com.google.cloud.Tuple;
 import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
@@ -430,16 +432,39 @@ public InsertAllResponse insertAll(InsertAllRequest request) {
     requestPb.setIgnoreUnknownValues(request.ignoreUnknownValues());
     requestPb.setSkipInvalidRows(request.skipInvalidRows());
     requestPb.setTemplateSuffix(request.getTemplateSuffix());
+    // Using an array of size 1 here to have a mutable boolean variable, which can be modified in
+    // an anonymous inner class.
+    final boolean[] allInsertIdsSet = {true};
     List<Rows> rowsPb = Lists.transform(request.getRows(), new Function<RowToInsert, Rows>() {
       @Override
       public Rows apply(RowToInsert rowToInsert) {
+        allInsertIdsSet[0] &= rowToInsert.getId() != null;
         return new Rows().setInsertId(rowToInsert.getId()).setJson(rowToInsert.getContent());
       }
     });
     requestPb.setRows(rowsPb);
-    return InsertAllResponse.fromPb(
-        bigQueryRpc.insertAll(tableId.getProject(), tableId.getDataset(), tableId.getTable(),
-            requestPb));
+
+    TableDataInsertAllResponse responsePb;
+    if (allInsertIdsSet[0]) {
+      // allowing retries only if all row insertIds are set (used for deduplication)
+      try {
+        responsePb = runWithRetries(
+            new Callable<TableDataInsertAllResponse>() {
+              @Override
+              public TableDataInsertAllResponse call() throws Exception {
+                return bigQueryRpc.insertAll(tableId.getProject(), tableId.getDataset(),
+                    tableId.getTable(), requestPb);
+              }
+            }, getOptions().getRetrySettings(), EXCEPTION_HANDLER, getOptions().getClock());
+      } catch (RetryHelperException e) {
+        throw BigQueryException.translateAndThrow(e);
+      }
+    } else {
+      responsePb = bigQueryRpc.insertAll(tableId.getProject(), tableId.getDataset(),
+          tableId.getTable(), requestPb);
+    }
+
+    return InsertAllResponse.fromPb(responsePb);
   }
 
   @Override

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -80,6 +80,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Sets;
 import com.google.common.io.BaseEncoding;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -1093,8 +1094,11 @@ public void testExtractJob() throws InterruptedException, TimeoutException {
     Job remoteExtractJob = bigquery.create(JobInfo.of(extractConfiguration));
     remoteExtractJob = remoteExtractJob.waitFor();
     assertNull(remoteExtractJob.getStatus().getError());
-    assertEquals(CSV_CONTENT,
-        new String(storage.readAllBytes(BUCKET, EXTRACT_FILE), StandardCharsets.UTF_8));
+
+    String extractedCsv =
+        new String(storage.readAllBytes(BUCKET, EXTRACT_FILE), StandardCharsets.UTF_8);
+    assertEquals(
+        Sets.newHashSet(CSV_CONTENT.split("\n")), Sets.newHashSet(extractedCsv.split("\n")));
     assertTrue(bigquery.delete(DATASET, tableName));
   }