From 8e70f89f502785725c30f111ca8bccaa45cb4d2a Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 15 Oct 2025 01:13:10 +0000 Subject: [PATCH] Add max retries override + unit test Signed-off-by: Ubuntu --- .../com/databricks/sdk/core/ApiClient.java | 9 +++++- .../databricks/sdk/core/DatabricksConfig.java | 13 ++++++++ .../databricks/sdk/core/ApiClientTest.java | 30 +++++++++++++++++++ .../sdk/core/DatabricksConfigTest.java | 17 +++++++++++ 4 files changed, 68 insertions(+), 1 deletion(-) 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..d6c691ad9 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 @@ -37,6 +37,7 @@ public static class Builder { private Function getHostFunc; private Function 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); 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..8d145b8ed 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 @@ -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; } 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..2776d1cc1 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 @@ -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(); 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..081dc42c5 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 testMaxRetriesSetAndGet() { + DatabricksConfig config = new DatabricksConfig().setMaxRetries(10); + assertEquals(Integer.valueOf(10), config.getMaxRetries()); + } + + @Test + public void testMaxRetriesEnvironmentVariable() { + Map 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()); + } }