[PECOBLR-1219] Added Api retriable codes URL param

NEXT_CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Added
 - Added the EnableTokenFederation url param to enable or disable Token federation feature. By default it is set to 1
+- Added the ApiRetriableHttpCodes, ApiRetryTimeout url params to enable retries for specific HTTP codes irrespective of Retry-After header. By default the HTTP codes list is empty.
 
 ### Updated
 - Added validation for positive integer configuration properties (RowsFetchedPerBlock, BatchInsertSize, etc.) to prevent hangs and errors when set to zero or negative values.

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java

@@ -4,6 +4,7 @@
 import static com.databricks.jdbc.common.DatabricksJdbcUrlParams.AUTH_SCOPE;
 import static com.databricks.jdbc.common.DatabricksJdbcUrlParams.DEFAULT_STRING_COLUMN_LENGTH;
 import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_ROW_LIMIT_PER_BLOCK;
+import static com.databricks.jdbc.common.util.StringUtil.parseIntegerSet;
 import static com.databricks.jdbc.common.util.UserAgentManager.USER_AGENT_SEA_CLIENT;
 import static com.databricks.jdbc.common.util.UserAgentManager.USER_AGENT_THRIFT_CLIENT;
 import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty;
@@ -28,6 +29,7 @@
 import com.google.common.collect.ImmutableMap;
 import java.net.URI;
 import java.util.*;
+import java.util.Collections;
 import java.util.regex.Matcher;
 import java.util.stream.Collectors;
 import org.apache.http.client.utils.URIBuilder;
@@ -672,6 +674,17 @@ public int getRateLimitRetryTimeout() {
     return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.RATE_LIMIT_RETRY_TIMEOUT));
   }
 
+  @Override
+  public Set<Integer> getApiRetriableHttpCodes() {
+    String codes = getParameter(DatabricksJdbcUrlParams.API_RETRIABLE_HTTP_CODES);
+    return parseIntegerSet(codes);
+  }
+
+  @Override
+  public int getApiRetryTimeout() {
+    return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.API_RETRY_TIMEOUT));
+  }
+
   @Override
   public int getIdleHttpConnectionExpiry() {
     return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.IDLE_HTTP_CONNECTION_EXPIRY));

src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java

@@ -7,6 +7,7 @@
 import com.databricks.sdk.core.utils.Cloud;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public interface IDatabricksConnectionContext {
 
@@ -172,6 +173,10 @@ public interface IDatabricksConnectionContext {
 
   int getRateLimitRetryTimeout();
 
+  Set<Integer> getApiRetriableHttpCodes();
+
+  int getApiRetryTimeout();
+
   int getIdleHttpConnectionExpiry();
 
   boolean supportManyParameters();

src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java

@@ -187,7 +187,13 @@ public enum DatabricksJdbcUrlParams {
       "Disable requesting OAuth refresh tokens (omit offline_access unless explicitly provided)",
       "1"),
   ENABLE_TOKEN_FEDERATION(
-      "EnableTokenFederation", "Enable token federation for authentication", "1");
+      "EnableTokenFederation", "Enable token federation for authentication", "1"),
+  API_RETRIABLE_HTTP_CODES(
+      "ApiRetriableHttpCodes",
+      "Comma-separated list of HTTP status codes that should be retried irrespective of Retry-After header.",
+      ""),
+  API_RETRY_TIMEOUT(
+      "ApiRetryTimeout", "Timeout for retrying API retriable codes in seconds", "120");
 
   private final String paramName;
   private final String defaultValue;

src/main/java/com/databricks/jdbc/common/util/StringUtil.java

@@ -1,7 +1,10 @@
 package com.databricks.jdbc.common.util;
 
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 public class StringUtil {
   public static String convertJdbcEscapeSequences(String sql) {
@@ -88,4 +91,23 @@ public static String removeRedundantEscapeClause(String sql) {
     // followed by delimiter (whitespace, semicolon, end-of-string, closing paren, or comma)
     return sql.replaceAll("(?i)\\s+ESCAPE\\s+['\"]['\"](?=\\s|;|$|\\)|,)", "");
   }
+
+  /**
+   * Parses a comma-separated string of integers into a Set.
+   *
+   * @param input comma-separated string of integers (e.g., "500, 503, 504")
+   * @return Set of parsed integers, or empty set if input is null/empty or all values are invalid
+   * @throws NumberFormatException if any value cannot be parsed as an integer (after validation)
+   */
+  public static Set<Integer> parseIntegerSet(String input) {
+    if (input.trim().isEmpty()) {
+      return Collections.emptySet();
+    }
+    return Arrays.stream(input.split(","))
+        .map(String::trim)
+        .filter(s -> !s.isEmpty())
+        .filter(num -> num.matches("\\d+")) // Ensure only positive integers
+        .map(Integer::parseInt)
+        .collect(Collectors.toSet());
+  }
 }

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java

@@ -9,6 +9,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.util.Objects;
+import java.util.Set;
 import org.apache.http.HttpResponse;
 import org.apache.http.HttpResponseInterceptor;
 import org.apache.http.HttpStatus;
@@ -28,6 +29,7 @@ public class DatabricksHttpRetryHandler
   private static final String TEMP_UNAVAILABLE_ACCUMULATED_TIME_KEY =
       "tempUnavailableAccumulatedTime";
   private static final String RATE_LIMIT_ACCUMULATED_TIME_KEY = "rateLimitAccumulatedTime";
+  private static final String API_CODES_ACCUMULATED_TIME_KEY = "apiCodesAccumulatedTime";
   static final String RETRY_AFTER_HEADER = "Retry-After";
   private static final int DEFAULT_BACKOFF_FACTOR = 2; // Exponential factor
   private static final int MIN_BACKOFF_INTERVAL = 1000; // 1s
@@ -134,10 +136,12 @@ public boolean retryRequest(IOException exception, int executionCount, HttpConte
 
     // check if retry interval is valid for 503 and 429
     int retryInterval = (int) context.getAttribute(RETRY_INTERVAL_KEY);
-    if ((statusCode == HttpStatus.SC_SERVICE_UNAVAILABLE
-            || statusCode == HttpStatus.SC_TOO_MANY_REQUESTS)
-        && retryInterval == -1) {
+
+    Set<Integer> apiRetriableCodes = connectionContext.getApiRetriableHttpCodes();
+    boolean isInCustomRetriableCodes = apiRetriableCodes.contains(statusCode);
+    if (retryInterval == -1 && !isInCustomRetriableCodes) {
       // This case arises when the server does not send the retryAfter header
+      // and the status code is not in the custom retriable codes list
       LOGGER.warn(
           "Invalid retry interval in the context "
               + context
@@ -146,9 +150,15 @@ public boolean retryRequest(IOException exception, int executionCount, HttpConte
       return false;
     }
 
+    // If no retry-after header (retryInterval == -1), calculate delay using exponential backoff
+    if (retryInterval == -1) {
+      retryInterval = (int) (calculateExponentialBackoff(executionCount) / 1000);
+    }
+
     long tempUnavailableAccumulatedTime =
         getAccumulatedTime(context, TEMP_UNAVAILABLE_ACCUMULATED_TIME_KEY);
     long rateLimitAccumulatedTime = getAccumulatedTime(context, RATE_LIMIT_ACCUMULATED_TIME_KEY);
+    long apiCodesAccumulatedTime = getAccumulatedTime(context, API_CODES_ACCUMULATED_TIME_KEY);
 
     // check if retry timeout has been hit for error code 503
     if (statusCode == HttpStatus.SC_SERVICE_UNAVAILABLE
@@ -174,6 +184,19 @@ public boolean retryRequest(IOException exception, int executionCount, HttpConte
       return false;
     }
 
+    // check if retry timeout has been hit for custom API retriable codes
+    if (apiRetriableCodes.contains(statusCode)
+        && statusCode != HttpStatus.SC_SERVICE_UNAVAILABLE
+        && statusCode != HttpStatus.SC_TOO_MANY_REQUESTS
+        && apiCodesAccumulatedTime + retryInterval > connectionContext.getApiRetryTimeout()) {
+      LOGGER.warn(
+          "ApiRetry timeout "
+              + connectionContext.getApiRetryTimeout()
+              + " has been hit for the error: "
+              + exception.getMessage());
+      return false;
+    }
+
     // check if request method is retryable
     boolean isRequestMethodRetryable =
         isRequestMethodRetryable(
@@ -190,6 +213,8 @@ public boolean retryRequest(IOException exception, int executionCount, HttpConte
     } else if (statusCode == HttpStatus.SC_TOO_MANY_REQUESTS) {
       context.setAttribute(
           RATE_LIMIT_ACCUMULATED_TIME_KEY, rateLimitAccumulatedTime + retryInterval);
+    } else if (apiRetriableCodes.contains(statusCode)) {
+      context.setAttribute(API_CODES_ACCUMULATED_TIME_KEY, apiCodesAccumulatedTime + retryInterval);
     }
 
     // calculate the delay and sleep for that duration
@@ -237,6 +262,9 @@ private static void initializeRetryAccumulatedTimeIfNotExist(HttpContext httpCon
     if (httpContext.getAttribute(RATE_LIMIT_ACCUMULATED_TIME_KEY) == null) {
       httpContext.setAttribute(RATE_LIMIT_ACCUMULATED_TIME_KEY, 0L);
     }
+    if (httpContext.getAttribute(API_CODES_ACCUMULATED_TIME_KEY) == null) {
+      httpContext.setAttribute(API_CODES_ACCUMULATED_TIME_KEY, 0L);
+    }
   }
 
   private static long getAccumulatedTime(HttpContext context, String key) {
@@ -256,6 +284,11 @@ protected void doSleepForDelay(long delayMillis) {
 
   /** Check if the request is retryable based on the status code and any connection preferences. */
   private boolean isStatusCodeRetryable(int statusCode) {
+    Set<Integer> apiRetriableCodes = connectionContext.getApiRetriableHttpCodes();
+    if (apiRetriableCodes.contains(statusCode)) {
+      return true;
+    }
+
     switch (statusCode) {
       case HttpStatus.SC_SERVICE_UNAVAILABLE:
         return connectionContext.shouldRetryTemporarilyUnavailableError();

src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java

@@ -1,7 +1,9 @@
 package com.databricks.jdbc.common.util;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.Set;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.junit.jupiter.MockitoExtension;
@@ -86,4 +88,18 @@ public void testPreserveCustomEscapeClause() {
     String sql = "SELECT * FROM table WHERE name LIKE 'pattern#%' ESCAPE '#'";
     assertEquals(sql, StringUtil.removeRedundantEscapeClause(sql));
   }
+
+  @Test
+  public void testParseIntegerSet() {
+    // Valid input
+    Set<Integer> result = StringUtil.parseIntegerSet("500,503,504");
+    assertEquals(3, result.size());
+    assertTrue(result.contains(500));
+    assertTrue(result.contains(503));
+    assertTrue(result.contains(504));
+
+    // Empty or whitespace
+    assertTrue(StringUtil.parseIntegerSet("").isEmpty());
+    assertTrue(StringUtil.parseIntegerSet("   ").isEmpty());
+  }
 }

src/test/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandlerTest.java

@@ -264,6 +264,32 @@ void testRetryRequestWithInvalidRetryInterval() {
     assertFalse(shouldRetry429, "Should return false when retryInterval is -1 for status 429");
   }
 
+  @Test
+  void testApiRetriableCodesWithoutRetryAfterHeader() throws IOException {
+    when(mockConnectionContext.getApiRetriableHttpCodes())
+        .thenReturn(java.util.Set.of(HttpStatus.SC_INTERNAL_SERVER_ERROR));
+    when(mockConnectionContext.getApiRetryTimeout()).thenReturn(120);
+
+    HttpRequest request = createRequest("GET", "/api/data");
+    httpContext.setAttribute(HttpCoreContext.HTTP_REQUEST, request);
+
+    // 503 WITHOUT Retry-After header - should use exponential backoff
+    HttpResponse response = createResponse(HttpStatus.SC_INTERNAL_SERVER_ERROR);
+    assertThrows(
+        DatabricksRetryHandlerException.class, () -> retryHandler.process(response, httpContext));
+
+    assertTrue(
+        retryHandler.retryRequest(
+            new DatabricksRetryHandlerException("Test", HttpStatus.SC_INTERNAL_SERVER_ERROR),
+            1,
+            httpContext));
+
+    // Verify exponential backoff delay was used
+    long expectedDelay = DatabricksHttpRetryHandler.calculateExponentialBackoff(1);
+    assertEquals(1, sleepDurations.size());
+    assertEquals(expectedDelay, sleepDurations.get(0));
+  }
+
   private HttpResponse createResponse(int statusCode) {
     return createResponse(statusCode, null);
   }