feat: Updated BigQueryRetryAlgorithm so that it can retry on RateLimit Errors using RegEx (#1499)

* Modified RATE_LIMIT_ERROR_MSG and added testRateLimitRegEx for testing regex matching

* Added JOB_RATE_LIMIT_EXCEEDED_MSG and RATE_LIMIT_EXCEEDED_REGEX pattern

* Implemented retryOnRegEx to include retry on RegEx patterns

* Updated DEFAULT_RETRY_CONFIG to include retry on JOB_RATE_LIMIT_EXCEEDED_MSG and RATE_LIMIT_EXCEEDED_REGEX

* Modified shouldRetryBasedOnBigQueryRetryConfig method to try RegEx matching if all the contains checks fail

Author: prash-mi

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

@@ -19,4 +19,9 @@
 public class BigQueryErrorMessages {
   public static final String RATE_LIMIT_EXCEEDED_MSG =
       "Exceeded rate limits:"; // Error Message for RateLimitExceeded Error
+  public static final String JOB_RATE_LIMIT_EXCEEDED_MSG = "Job exceeded rate limits:";
+
+  public class RetryRegExPatterns {
+    public static final String RATE_LIMIT_EXCEEDED_REGEX = ".*exceed.*rate.*limit.*";
+  }
 }

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

@@ -240,7 +240,9 @@ public Page<FieldValueList> getNextPage() {
   private static final BigQueryRetryConfig DEFAULT_RETRY_CONFIG =
       BigQueryRetryConfig.newBuilder()
           .retryOnMessage(BigQueryErrorMessages.RATE_LIMIT_EXCEEDED_MSG)
-          .build(); // retry config with Error Message for RateLimitExceeded Error
+          .retryOnMessage(BigQueryErrorMessages.JOB_RATE_LIMIT_EXCEEDED_MSG)
+          .retryOnRegEx(BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX)
+          .build(); // retry config with Error Messages and RegEx for RateLimitExceeded Error
 
   BigQueryImpl(BigQueryOptions options) {
     super(options);

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

@@ -27,6 +27,7 @@
 import com.google.api.gax.retrying.TimedRetryAlgorithmWithContext;
 import java.util.Iterator;
 import java.util.concurrent.CancellationException;
+import java.util.regex.Pattern;
 
 public class BigQueryRetryAlgorithm<ResponseT> extends RetryAlgorithm<ResponseT> {
   private final BigQueryRetryConfig bigQueryRetryConfig;
@@ -69,17 +70,35 @@ private boolean shouldRetryBasedOnBigQueryRetryConfig(
      */
     String errorDesc;
     if (previousThrowable != null && (errorDesc = previousThrowable.getMessage()) != null) {
+      errorDesc = errorDesc.toLowerCase(); // for case insensitive comparison
       for (Iterator<String> retriableMessages =
               bigQueryRetryConfig.getRetriableErrorMessages().iterator();
           retriableMessages.hasNext(); ) {
-        if (errorDesc.contains(retriableMessages.next())) { // Error message should be retried
+        if (errorDesc.contains(
+            retriableMessages
+                .next()
+                .toLowerCase())) { // Error message should be retried, implementing cases
+          // insensitive match
+          return true;
+        }
+      }
+      // Check if there's a regex which matches the error message. This avoids too many regex
+      // matches which is expensive
+      for (Iterator<String> retriableRegExes = bigQueryRetryConfig.getRetriableRegExes().iterator();
+          retriableRegExes.hasNext(); ) {
+        if (matchRegEx(retriableRegExes.next(), errorDesc)) {
           return true;
         }
       }
     }
     return false;
   }
 
+  public static boolean matchRegEx(
+      String retriableRegEx, String errorDesc) { // cases insensitive match regex matching
+    return Pattern.matches(retriableRegEx.toLowerCase(), errorDesc.toLowerCase());
+  }
+
   /*Duplicating this method as it can not be inherited from the RetryAlgorithm due to the default access modifier*/
   boolean shouldRetryBasedOnResult(
       RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) {

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

@@ -21,18 +21,25 @@
 
 public class BigQueryRetryConfig {
   private final ImmutableSet<String> retriableErrorMessages;
+  private final ImmutableSet<String> retriableRegExes;
 
   private BigQueryRetryConfig(Builder builder) {
     retriableErrorMessages = builder.retriableErrorMessages.build();
+    retriableRegExes = builder.retriableRegExes.build();
   }
 
   public ImmutableSet<String> getRetriableErrorMessages() {
     return retriableErrorMessages;
   }
 
+  public ImmutableSet<String> getRetriableRegExes() {
+    return retriableRegExes;
+  }
+
   // BigQueryRetryConfig builder
   public static class Builder {
     private final ImmutableSet.Builder<String> retriableErrorMessages = ImmutableSet.builder();
+    private final ImmutableSet.Builder<String> retriableRegExes = ImmutableSet.builder();
 
     private Builder() {}
 
@@ -43,6 +50,13 @@ public final Builder retryOnMessage(String... errorMessages) {
       return this;
     }
 
+    public final Builder retryOnRegEx(String... regExPatterns) {
+      for (String regExPattern : regExPatterns) {
+        retriableRegExes.add(checkNotNull(regExPattern));
+      }
+      return this;
+    }
+
     public BigQueryRetryConfig build() {
       return new BigQueryRetryConfig(this);
     }

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -508,11 +508,12 @@ public class BigQueryImplTest {
           .setEtag(ETAG)
           .setVersion(1)
           .build();
-
   private BigQueryOptions options;
   private BigQueryRpcFactory rpcFactoryMock;
   private BigQueryRpc bigqueryRpcMock;
   private BigQuery bigquery;
+  private static final String RATE_LIMIT_ERROR_MSG =
+      "Job exceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";
 
   @Captor private ArgumentCaptor<Map<BigQueryRpc.Option, Object>> capturedOptions;
   @Captor private ArgumentCaptor<com.google.api.services.bigquery.model.Job> jobCapture;
@@ -2439,9 +2440,7 @@ public void testFastQueryRateLimitIdempotency() throws Exception {
         .thenThrow(new BigQueryException(504, "Gateway Timeout"))
         .thenThrow(
             new BigQueryException(
-                400,
-                BigQueryErrorMessages
-                    .RATE_LIMIT_EXCEEDED_MSG)) // retrial on based on RATE_LIMIT_EXCEEDED_MSG
+                400, RATE_LIMIT_ERROR_MSG)) // retrial on based on RATE_LIMIT_EXCEEDED_MSG
         .thenReturn(responsePb);
 
     bigquery =
@@ -2470,6 +2469,30 @@ public void testFastQueryRateLimitIdempotency() throws Exception {
     verify(bigqueryRpcMock, times(6)).queryRpc(eq(PROJECT), requestPbCapture.capture());
   }
 
+  @Test
+  public void testRateLimitRegEx() throws Exception {
+    String msg2 =
+        "Job eceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";
+    String msg3 = "exceeded rate exceeded quota for table update";
+    String msg4 = "exceeded rate limits";
+    assertTrue(
+        BigQueryRetryAlgorithm.matchRegEx(
+            BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
+            RATE_LIMIT_ERROR_MSG));
+    assertFalse(
+        BigQueryRetryAlgorithm.matchRegEx(
+            BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
+            msg2.toLowerCase()));
+    assertFalse(
+        BigQueryRetryAlgorithm.matchRegEx(
+            BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
+            msg3.toLowerCase()));
+    assertTrue(
+        BigQueryRetryAlgorithm.matchRegEx(
+            BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
+            msg4.toLowerCase()));
+  }
+
   @Test
   public void testFastQueryDDLShouldRetry() throws Exception {
     com.google.api.services.bigquery.model.QueryResponse responsePb =