diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 0fd0b1445..6ba141908 100644 --- a/NEXT_CHANGELOG.md +++ b/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 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ApiClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ApiClient.java index 2d4eeadc0..20ebad038 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ApiClient.java +++ b/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; } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 775b52a79..68cb0fff0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/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; } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/NoRetryStrategy.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/NoRetryStrategy.java new file mode 100644 index 000000000..5b72d06bf --- /dev/null +++ b/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; + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/NoRetryStrategyPicker.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/NoRetryStrategyPicker.java new file mode 100644 index 000000000..b87e018cc --- /dev/null +++ b/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; + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/ApiClientTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/ApiClientTest.java index efd4fbae2..4cfe0edaf 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/ApiClientTest.java +++ b/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(); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 88a466a32..0f3fecf6d 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/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 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()); + } }