Adding retry for concurrent merge failed serialization

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/MergeQueries.java

@@ -20,22 +20,28 @@
 package com.wepay.kafka.connect.bigquery;
 
 import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.Field;
 import com.google.cloud.bigquery.FieldList;
 import com.google.cloud.bigquery.QueryJobConfiguration;
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.TableId;
 import com.google.common.annotations.VisibleForTesting;
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkTaskConfig;
+import com.wepay.kafka.connect.bigquery.exception.BigQueryConnectException;
 import com.wepay.kafka.connect.bigquery.exception.ExpectedInterruptException;
+import com.wepay.kafka.connect.bigquery.utils.SleepUtils;
 import com.wepay.kafka.connect.bigquery.write.batch.KCBQThreadPoolExecutor;
 import com.wepay.kafka.connect.bigquery.write.batch.MergeBatches;
+import com.wepay.kafka.connect.bigquery.write.row.BigQueryErrorResponses;
 import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.errors.RetriableException;
 import org.apache.kafka.connect.sink.SinkTaskContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -139,7 +145,26 @@ private void mergeFlush(
           batchNumber, intTable(intermediateTable));
       String mergeFlushQuery = mergeFlushQuery(intermediateTable, destinationTable, batchNumber);
       logger.trace(mergeFlushQuery);
-      bigQuery.query(QueryJobConfiguration.of(mergeFlushQuery));
+
+      int attempt = 0;
+      boolean success = false;
+      while (!success) {
+        try {
+          bigQuery.query(QueryJobConfiguration.of(mergeFlushQuery));
+          success = true;
+        } catch (BigQueryException e) {
+          if (BigQueryErrorResponses.isCouldNotSerializeAccessError(e)) {
+            attempt++;
+            if (attempt == 30) {
+              throw new BigQueryConnectException("Failed to merge rows to destination table `" + destinationTable + "` within " + attempt 
+                                                  + "  attempts due to BQ write serialization error.", e);
+            }
+            SleepUtils.waitRandomTime(10000, 20000);
+          } else {
+            throw e;
+          }
+        }
+      }
       logger.trace("Merge from {} to {} completed",
           intTable(intermediateTable), destTable(destinationTable));
 

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/utils/SleepUtils.java

@@ -0,0 +1,10 @@
+package com.wepay.kafka.connect.bigquery.utils;
+
+import java.util.concurrent.ThreadLocalRandom;
+
+public final class SleepUtils {
+
+    public static void waitRandomTime(int sleepMs, int jitterMs) throws InterruptedException {
+        Thread.sleep(sleepMs + ThreadLocalRandom.current().nextInt(jitterMs));
+    }
+}

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/write/row/BigQueryErrorResponses.java

@@ -43,6 +43,7 @@ public class BigQueryErrorResponses {
 
   private static final String BAD_REQUEST_REASON = "badRequest";
   private static final String INVALID_REASON = "invalid";
+  private static final String INVALID_QUERY_REASON = "invalidQuery";
   private static final String NOT_FOUND_REASON = "notFound";
   private static final String QUOTA_EXCEEDED_REASON = "quotaExceeded";
   private static final String RATE_LIMIT_EXCEEDED_REASON = "rateLimitExceeded";
@@ -113,6 +114,13 @@ public static boolean isIOError(BigQueryException error) {
     return BigQueryException.UNKNOWN_CODE == error.getCode()
         && error.getCause() instanceof IOException;
   }
+
+  public static boolean isCouldNotSerializeAccessError(BigQueryException exception) {
+    return BAD_REQUEST_CODE == exception.getCode()
+            && INVALID_QUERY_REASON.equals(exception.getReason())
+            && message(exception.getError()).startsWith("Could not serialize access to");
+  }
+
   /**
    * Returns whether the error code and the description string match to authentication errors.
    * See also <a href="https://cloud.google.com/bigquery/docs/error-messages#autherrors">here</a>.
@@ -125,6 +133,7 @@ public static boolean isAuthenticationError(BigQueryException error) {
             ||
             err.contains(String.valueOf(AUTHENTICATION_ERROR_CODE));
   }
+
   public static boolean isUnrecognizedFieldError(BigQueryError error) {
     return INVALID_REASON.equals(reason(error))
         && message(error).startsWith("no such field: ");