Add NoRetryStrategy

Author: renaudhartert-db

NEXT_CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### New Features and Improvements
 
+* Add a new config attribute `DATABRICKS_DISABLE_RETRIES` to disable the 
+  default retry mechanism.
+
 ### Bug Fixes
 
 ### Documentation

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

@@ -6,6 +6,7 @@
 import com.databricks.sdk.core.http.Request;
 import com.databricks.sdk.core.http.RequestOptions;
 import com.databricks.sdk.core.http.Response;
+import com.databricks.sdk.core.retry.NoRetryStrategyPicker;
 import com.databricks.sdk.core.retry.RequestBasedRetryStrategyPicker;
 import com.databricks.sdk.core.retry.RetryStrategy;
 import com.databricks.sdk.core.retry.RetryStrategyPicker;
@@ -49,8 +50,14 @@ public Builder withDatabricksConfig(DatabricksConfig config) {
       this.httpClient = config.getHttpClient();
       this.debugTruncateBytes = config.getDebugTruncateBytes();
       this.accountId = config.getAccountId();
-      this.retryStrategyPicker = new RequestBasedRetryStrategyPicker(config.getHost());
       this.isDebugHeaders = config.isDebugHeaders();
+
+      if (config.getDisableRetries()) {
+        this.retryStrategyPicker = new NoRetryStrategyPicker();
+      } else {
+        this.retryStrategyPicker = new RequestBasedRetryStrategyPicker(config.getHost());
+      }
+
       return this;
     }
 

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

@@ -164,6 +164,10 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")
   private Boolean disableAsyncTokenRefresh;
 
+  /** Disable retries by default when set to true. */
+  @ConfigAttribute(env = "DATABRICKS_DISABLE_RETRIES")
+  private Boolean disableRetries;
+
   /**
    * The duration to wait for a browser response during U2M authentication before timing out. If set
    * to 0 or null, the connector waits for an indefinite amount of time.
@@ -609,6 +613,15 @@ public DatabricksConfig setDisableAsyncTokenRefresh(boolean disableAsyncTokenRef
     return this;
   }
 
+  public boolean getDisableRetries() {
+    return disableRetries != null && disableRetries;
+  }
+
+  public DatabricksConfig setDisableRetries(boolean disableRetries) {
+    this.disableRetries = disableRetries;
+    return this;
+  }
+
   public Duration getOAuthBrowserAuthTimeout() {
     return oauthBrowserAuthTimeout;
   }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/NoRetryStrategy.java

@@ -0,0 +1,15 @@
+package com.databricks.sdk.core.retry;
+
+import com.databricks.sdk.core.DatabricksError;
+
+/**
+ * This class implements a retry strategy that never retries any requests. All requests are
+ * considered non-retriable regardless of the error type or status code.
+ */
+public class NoRetryStrategy implements RetryStrategy {
+
+  @Override
+  public boolean isRetriable(DatabricksError databricksError) {
+    return false;
+  }
+}

databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/NoRetryStrategyPicker.java

@@ -0,0 +1,22 @@
+package com.databricks.sdk.core.retry;
+
+import com.databricks.sdk.core.http.Request;
+
+/**
+ * A RetryStrategyPicker that always returns a NoRetryStrategy, effectively disabling retries for
+ * all requests regardless of their type or characteristics.
+ */
+public class NoRetryStrategyPicker implements RetryStrategyPicker {
+  private static final NoRetryStrategy NO_RETRY_STRATEGY = new NoRetryStrategy();
+
+  /**
+   * Returns a NoRetryStrategy for any given request, effectively disabling retries.
+   *
+   * @param request The request for which the retry strategy is needed (ignored).
+   * @return A NoRetryStrategy instance that will never retry any request.
+   */
+  @Override
+  public RetryStrategy getRetryStrategy(Request request) {
+    return NO_RETRY_STRATEGY;
+  }
+}

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

@@ -410,6 +410,40 @@ public HeaderFactory configure(DatabricksConfig config) {
     }
   }
 
+  @Test
+  void verifyNoRetryWhenRetriesDisabled() throws IOException {
+    Request req = getBasicRequest();
+    DatabricksConfig config =
+        new DatabricksConfig()
+            .setHost("http://my.host")
+            .setDisableRetries(true)
+            .setCredentialsProvider(new DummyCredentialsProvider());
+    ApiClient client =
+        getApiClient(
+            config, req, Arrays.asList(getTooManyRequestsResponse(req), getSuccessResponse(req)));
+
+    DatabricksError exception =
+        assertThrows(
+            DatabricksError.class,
+            () ->
+                client.execute(
+                    new Request("GET", req.getUri().getPath()), MyEndpointResponse.class));
+
+    assertEquals("TOO_MANY_REQUESTS", exception.getErrorCode());
+    assertEquals(429, exception.getStatusCode());
+  }
+
+  @Test
+  void verifyRetriesWorkWhenEnabled() throws IOException {
+    Request req = getBasicRequest();
+    // Verify that the client retries by default.
+    runApiClientTest(
+        req,
+        Arrays.asList(getTooManyRequestsResponse(req), getSuccessResponse(req)),
+        MyEndpointResponse.class,
+        new MyEndpointResponse().setKey("value")); // should succeed after retry
+  }
+
   @Test
   void populateHostFromCredentialProvider() throws IOException {
     Request req = getBasicRequest();

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 testDisableRetriesSetAndGet() {
+    DatabricksConfig config = new DatabricksConfig().setDisableRetries(true);
+    assertEquals(true, config.getDisableRetries());
+  }
+
+  @Test
+  public void testDisableRetriesEnvironmentVariable() {
+    Map<String, String> env = new HashMap<>();
+    env.put("DATABRICKS_DISABLE_RETRIES", "true");
+
+    DatabricksConfig config = new DatabricksConfig();
+    config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
+
+    assertEquals(true, config.getDisableRetries());
+  }
 }