From a37722faaf6518db5895546ace236782d44d23b3 Mon Sep 17 00:00:00 2001
From: Thomas Powell <tpowell@palantir.com>
Date: Wed, 23 Apr 2025 10:15:08 +0100
Subject: [PATCH] Capture DatabricksError when retrying API calls.

Signed-off-by: Thomas Powell <tpowell@palantir.com>
---
 NEXT_CHANGELOG.md                             |  1 +
 .../com/databricks/sdk/core/ApiClient.java    |  5 ++--
 .../databricks/sdk/core/ApiClientTest.java    | 30 ++++++++++++-------
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
index cc0f26322..bc13c6b4e 100644
--- a/NEXT_CHANGELOG.md
+++ b/NEXT_CHANGELOG.md
@@ -9,5 +9,6 @@
 ### Documentation
 
 ### Internal Changes
+* Capture DatabricksError when retrying API calls ([#427](https://github.com/databricks/databricks-sdk-java/pull/427)).
 
 ### API Changes
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 25aeb13f1..a45590b4d 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
@@ -247,12 +247,13 @@ private Response executeInner(Request in, String path) {
       }
       if (attemptNumber == maxAttempts) {
         throw new DatabricksException(
-            String.format("Request %s failed after %d retries", in, maxAttempts), err);
+            String.format("Request %s failed after %d retries", in, maxAttempts), databricksError);
       }
 
       // Retry after a backoff.
       long sleepMillis = getBackoffMillis(out, attemptNumber);
-      LOG.debug(String.format("Retry %s in %dms", in.getRequestLine(), sleepMillis));
+      LOG.debug(
+          String.format("Retry %s in %dms", in.getRequestLine(), sleepMillis), databricksError);
       try {
         timer.sleep(sleepMillis);
       } catch (InterruptedException ex) {
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 b4c22b178..bad3dfadc 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
@@ -11,6 +11,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.IOException;
 import java.net.MalformedURLException;
 import java.net.URL;
@@ -88,12 +89,14 @@ private <T> void runApiClientTest(
     runApiClientTest(client, request, clazz, expectedResponse);
   }
 
-  private void runFailingApiClientTest(
+  @CanIgnoreReturnValue
+  private DatabricksException runFailingApiClientTest(
       Request request, List<ResponseProvider> responses, Class<?> clazz, String expectedMessage)
       throws IOException {
     DatabricksException exception =
         runFailingApiClientTest(request, responses, clazz, DatabricksException.class);
     assertEquals(exception.getMessage(), expectedMessage);
+    return exception;
   }
 
   private <T extends Throwable> T runFailingApiClientTest(
@@ -217,16 +220,21 @@ void retry429() throws IOException {
   @Test
   void failAfterTooManyRetries() throws IOException {
     Request req = getBasicRequest();
-    runFailingApiClientTest(
-        req,
-        Arrays.asList(
-            getTooManyRequestsResponseWithRetryAfterDateHeader(req),
-            getTooManyRequestsResponse(req),
-            getTooManyRequestsResponse(req),
-            getTooManyRequestsResponse(req),
-            getSuccessResponse(req)),
-        MyEndpointResponse.class,
-        "Request GET /api/my/endpoint failed after 4 retries");
+    DatabricksException exception =
+        runFailingApiClientTest(
+            req,
+            Arrays.asList(
+                getTooManyRequestsResponseWithRetryAfterDateHeader(req),
+                getTooManyRequestsResponse(req),
+                getTooManyRequestsResponse(req),
+                getTooManyRequestsResponse(req),
+                getSuccessResponse(req)),
+            MyEndpointResponse.class,
+            "Request GET /api/my/endpoint failed after 4 retries");
+    assertInstanceOf(DatabricksError.class, exception.getCause());
+    DatabricksError cause = (DatabricksError) exception.getCause();
+    assertEquals(cause.getErrorCode(), "TOO_MANY_REQUESTS");
+    assertEquals(cause.getStatusCode(), 429);
   }
 
   @Test