Add max retries override + unit test

Signed-off-by: Ubuntu <ubuntu@ip-172-30-2-35.ec2.internal>

Author: Ubuntu

databricks-sdk-java/src/main/java/com/databricks/sdk/core/ApiClient.java

@@ -37,6 +37,7 @@ public static class Builder {
     private Function<Void, String> getHostFunc;
     private Function<Void, String> getAuthTypeFunc;
     private Integer debugTruncateBytes;
+    private Integer maxRetries;
     private HttpClient httpClient;
     private String accountId;
     private RetryStrategyPicker retryStrategyPicker;
@@ -48,6 +49,7 @@ public Builder withDatabricksConfig(DatabricksConfig config) {
       this.getAuthTypeFunc = v -> config.getAuthType();
       this.httpClient = config.getHttpClient();
       this.debugTruncateBytes = config.getDebugTruncateBytes();
+      this.maxRetries = config.getMaxRetries();
       this.accountId = config.getAccountId();
       this.retryStrategyPicker = new RequestBasedRetryStrategyPicker(config.getHost());
       this.isDebugHeaders = config.isDebugHeaders();
@@ -84,6 +86,11 @@ public Builder withRetryStrategyPicker(RetryStrategyPicker retryStrategyPicker)
       return this;
     }
 
+    public Builder withMaxRetries(int maxRetries) {
+      this.maxRetries = maxRetries;
+      return this;
+    }
+
     public ApiClient build() {
       return new ApiClient(this);
     }
@@ -145,7 +152,7 @@ private ApiClient(Builder builder) {
       debugTruncateBytes = 96;
     }
 
-    maxAttempts = 4;
+    maxAttempts = builder.maxRetries != null ? builder.maxRetries : 4;
     mapper = SerDeUtils.createMapper();
     random = new Random();
     bodyLogger = new BodyLogger(mapper, 1024, debugTruncateBytes);

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -124,6 +124,10 @@ public class DatabricksConfig {
   /** Number of seconds for HTTP timeout */
   @ConfigAttribute() private Integer httpTimeoutSeconds;
 
+  /** Maximum number of retry attempts for failed requests. Default is 4. */
+  @ConfigAttribute(env = "DATABRICKS_MAX_RETRIES")
+  private Integer maxRetries;
+
   /** Truncate JSON fields in JSON above this limit. Default is 96. */
   @ConfigAttribute(env = "DATABRICKS_DEBUG_TRUNCATE_BYTES")
   private Integer debugTruncateBytes;
@@ -537,6 +541,15 @@ public DatabricksConfig setHttpTimeoutSeconds(int httpTimeoutSeconds) {
     return this;
   }
 
+  public Integer getMaxRetries() {
+    return maxRetries;
+  }
+
+  public DatabricksConfig setMaxRetries(int maxRetries) {
+    this.maxRetries = maxRetries;
+    return this;
+  }
+
   public Integer getDebugTruncateBytes() {
     return debugTruncateBytes;
   }

databricks-sdk-java/src/test/java/com/databricks/sdk/core/ApiClientTest.java

@@ -278,6 +278,36 @@ void failIdempotentRequestAfterTooManyRetries() throws IOException {
         "Request GET /api/2.0/sql/sessions/ failed after 4 retries");
   }
 
+  @Test
+  void failAfterConfigurableRetries() throws IOException {
+    Request req = getBasicRequest();
+
+    // Test with custom maxRetries=2
+    DatabricksConfig config =
+        new DatabricksConfig()
+            .setHost("http://my.host")
+            .setMaxRetries(2)
+            .setCredentialsProvider(new DummyCredentialsProvider());
+
+    ApiClient client =
+        getApiClient(
+            config,
+            req,
+            Arrays.asList(
+                getTooManyRequestsResponse(req),
+                getTooManyRequestsResponse(req),
+                getSuccessResponse(req))); // Would succeed on attempt 3, but maxRetries=2
+
+    DatabricksException exception =
+        assertThrows(
+            DatabricksException.class,
+            () ->
+                client.execute(
+                    new Request("GET", req.getUri().getPath()), MyEndpointResponse.class));
+
+    assertTrue(exception.getMessage().contains("failed after 2 retries"));
+  }
+
   @Test
   void testEmptyBody() throws IOException {
     MyEndpointResponse response = new MyEndpointResponse();

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

@@ -282,4 +282,21 @@ public void testEnvironmentVariableLoading() {
     assertEquals(Integer.valueOf(100), config.getDebugTruncateBytes());
     assertEquals(Integer.valueOf(50), config.getRateLimit());
   }
+
+  @Test
+  public void testMaxRetriesSetAndGet() {
+    DatabricksConfig config = new DatabricksConfig().setMaxRetries(10);
+    assertEquals(Integer.valueOf(10), config.getMaxRetries());
+  }
+
+  @Test
+  public void testMaxRetriesEnvironmentVariable() {
+    Map<String, String> env = new HashMap<>();
+    env.put("DATABRICKS_MAX_RETRIES", "15");
+
+    DatabricksConfig config = new DatabricksConfig();
+    config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
+
+    assertEquals(Integer.valueOf(15), config.getMaxRetries());
+  }
 }