Correctly return TooManyRequests when receiving HTTP 429 errors.

NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### Bug Fixes
 
+* Fix HTTP 429 (Too Many Requests) responses to correctly return `TooManyRequests` exception instead of generic `DatabricksError`.
+
 ### Security Vulnerabilities
 
 ### Documentation

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

@@ -221,9 +221,6 @@ private Response executeInner(Request in, String path, RequestOptions options) {
     while (true) {
       attemptNumber++;
 
-      IOException err = null;
-      Response out = null;
-
       // Authenticate the request. Failures should not be retried.
       in.withHeaders(authenticateFunc.apply(null));
 
@@ -239,27 +236,33 @@ private Response executeInner(Request in, String path, RequestOptions options) {
         userAgent += String.format(" auth/%s", authType);
       }
       in.withHeader("User-Agent", userAgent);
-
       options.applyOptions(in);
 
-      // Make the request, catching any exceptions, as we may want to retry.
+      Response response;
+      DatabricksError databricksError;
+
       try {
-        out = httpClient.execute(in);
+        response = httpClient.execute(in);
         if (LOG.isDebugEnabled()) {
-          LOG.debug(makeLogRecord(in, out));
+          LOG.debug(makeLogRecord(in, response));
         }
+
+        if (isResponseSuccessful(response)) {
+          return response; // stop here if the request succeeded
+        }
+
+        // The request did not succeed. Though, some errors are retriable and
+        // should be retried with exponential backoff.
+        databricksError = ApiErrors.getDatabricksError(response);
       } catch (IOException e) {
-        err = e;
         LOG.debug("Request {} failed", in, e);
+        // TODO: This is necesarry for backward compatibility as the code used
+        // to allow retries on IO errors. However, it is not clear if this is
+        // something we should continue to support.
+        databricksError = new DatabricksError("IO_ERROR", 523, e);
+        response = null;
       }
 
-      // Check if the request succeeded
-      if (isRequestSuccessful(out, err)) {
-        return out;
-      }
-      // The request did not succeed.
-      // Check if the request cannot be retried: if yes, retry after backoff, else throw the error.
-      DatabricksError databricksError = ApiErrors.getDatabricksError(out, err);
       if (!retryStrategy.isRetriable(databricksError)) {
         throw databricksError;
       }
@@ -269,7 +272,7 @@ private Response executeInner(Request in, String path, RequestOptions options) {
       }
 
       // Retry after a backoff.
-      long sleepMillis = getBackoffMillis(out, attemptNumber);
+      long sleepMillis = getBackoffMillis(response, attemptNumber);
       LOG.debug(
           String.format("Retry %s in %dms", in.getRequestLine(), sleepMillis), databricksError);
       try {
@@ -281,9 +284,8 @@ private Response executeInner(Request in, String path, RequestOptions options) {
     }
   }
 
-  private boolean isRequestSuccessful(Response response, Exception e) {
-    return e == null
-        && response.getStatusCode() >= 200
+  private boolean isResponseSuccessful(Response response) {
+    return response.getStatusCode() >= 200
         && response.getStatusCode() < 300
         && !PrivateLinkInfo.isPrivateLinkRedirect(response);
   }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/error/ApiErrors.java

@@ -5,8 +5,11 @@
 import com.databricks.sdk.core.error.details.ErrorDetails;
 import com.databricks.sdk.core.http.Response;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import java.io.*;
+import com.google.api.client.util.Strings;
+import java.io.IOException;
+import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import org.apache.commons.io.IOUtils;
@@ -16,73 +19,103 @@ public class ApiErrors {
   private static final ObjectMapper MAPPER = new ObjectMapper();
   private static final Pattern HTML_ERROR_REGEX = Pattern.compile("<pre>(.*)</pre>");
   private static final ErrorMapper ERROR_MAPPER = new ErrorMapper();
+  private static final String SCIM_NULL_DETAILS = "null";
 
-  public static DatabricksError getDatabricksError(Response out, Exception error) {
-    if (error != null) {
-      // If the endpoint did not respond to the request, interpret the exception.
-      return new DatabricksError("IO_ERROR", 523, error);
-    } else if (out.getStatusCode() == 429) {
-      return new DatabricksError("TOO_MANY_REQUESTS", "Current request has to be retried", 429);
+  // Empty error details used to normalize errors with no error details.
+  private static final ErrorDetails DEFAULT_ERROR_DETAILS = ErrorDetails.builder().build();
+
+  public static DatabricksError getDatabricksError(Response response) {
+    // Private link error handling depends purely on the response URL without
+    // looking at the error body.
+    if (PrivateLinkInfo.isPrivateLinkRedirect(response)) {
+      return ERROR_MAPPER.apply(response, new ApiErrorBody());
+    }
+
+    Optional<ApiErrorBody> optionalErrorBody = parseApiError(response);
+    if (!optionalErrorBody.isPresent()) {
+      // Purely infer the error based on its status code.
+      return ERROR_MAPPER.apply(response, new ApiErrorBody());
     }
 
-    ApiErrorBody errorBody = readErrorFromResponse(out);
-    return ERROR_MAPPER.apply(out, errorBody);
+    ApiErrorBody errorBody = optionalErrorBody.get();
+
+    // TODO: normalization should ideally happen at API call level,
+    // allowing each method to control its own normalization.
+    normalizeError(errorBody);
+
+    return ERROR_MAPPER.apply(response, errorBody);
   }
 
-  private static ApiErrorBody readErrorFromResponse(Response response) {
-    // Private link error handling depends purely on the response URL.
-    if (PrivateLinkInfo.isPrivateLinkRedirect(response)) {
-      return new ApiErrorBody();
+  private static void normalizeError(ApiErrorBody errorBody) {
+    // Guarantee that all errors have an ErrorDetails container even
+    // if it does not contains any specific error details.
+    if (errorBody.getErrorDetails() == null) {
+      errorBody.setErrorDetails(DEFAULT_ERROR_DETAILS);
     }
-    ApiErrorBody errorBody = parseApiError(response);
 
-    // Condense API v1.2 and SCIM error string and code into the message and errorCode fields of
-    // DatabricksApiException.
-    if (errorBody.getApi12Error() != null && !errorBody.getApi12Error().isEmpty()) {
+    // Error messages from older APIs used a different field to pass the
+    // error message. We keep this code for backward compatibility until
+    // these APIs are fully deprecated.
+    if (!Strings.isNullOrEmpty(errorBody.getApi12Error())) {
       errorBody.setMessage(errorBody.getApi12Error());
+      return;
+    }
+
+    if (!Strings.isNullOrEmpty(errorBody.getMessage())) {
+      return;
+    }
+
+    // SCIM error handling.
+    //
+    // TODO: This code is brittle and should likely be refactored to a more
+    // robust solution to detect SCIM errors. This will likely involve
+    // parsing the SCIM error at the API call level rather than normalizing.
+    if (Strings.isNullOrEmpty(errorBody.getScimDetail())) {
+      return;
+    }
+    if (SCIM_NULL_DETAILS.equals(errorBody.getScimDetail())) {
+      errorBody.setMessage("SCIM API Internal Error");
+    } else {
+      errorBody.setMessage(errorBody.getScimDetail());
     }
-    if (errorBody.getMessage() == null || errorBody.getMessage().isEmpty()) {
-      if (errorBody.getScimDetail() != null && !"null".equals(errorBody.getScimDetail())) {
-        errorBody.setMessage(errorBody.getScimDetail());
-      } else {
-        errorBody.setMessage("SCIM API Internal Error");
-      }
+    if (!Strings.isNullOrEmpty(errorBody.getScimType())) {
       String message = errorBody.getScimType() + " " + errorBody.getMessage();
       errorBody.setMessage(message.trim());
-      errorBody.setErrorCode("SCIM_" + errorBody.getScimStatus());
     }
-    if (errorBody.getErrorDetails() == null) {
-      errorBody.setErrorDetails(ErrorDetails.builder().build());
+    if (!Strings.isNullOrEmpty(errorBody.getScimStatus())) {
+      errorBody.setErrorCode("SCIM_" + errorBody.getScimStatus());
     }
-    return errorBody;
   }
 
   /**
    * The response is either a JSON response or a webpage. In the JSON case, it is parsed normally;
    * in the webpage case, the relevant error metadata is extracted from the response using a
    * heuristic-based approach.
    */
-  private static ApiErrorBody parseApiError(Response response) {
+  private static Optional<ApiErrorBody> parseApiError(Response response) {
+    InputStream in = response.getBody();
+    if (in == null) {
+      return Optional.empty();
+    }
+
+    // Read the body now, so we can try to parse as JSON and then
+    // fallback to old error handling logic.
+    String body;
     try {
-      InputStream in = response.getBody();
-      if (in == null) {
-        ApiErrorBody errorBody = new ApiErrorBody();
-        errorBody.setMessage(
-            String.format("Status response from server: %s", response.getStatus()));
-        return errorBody;
-      }
-
-      // Read the body now, so we can try to parse as JSON and then fallback to old error handling
-      // logic.
-      String body = IOUtils.toString(in, StandardCharsets.UTF_8);
-      try {
-        return MAPPER.readValue(body, ApiErrorBody.class);
-      } catch (IOException e) {
-        return parseUnknownError(response, body, e);
-      }
+      body = IOUtils.toString(in, StandardCharsets.UTF_8);
     } catch (IOException e) {
       throw new DatabricksException("Unable to read response body", e);
     }
+
+    if (Strings.isNullOrEmpty(body)) {
+      return Optional.empty();
+    }
+
+    try {
+      return Optional.of(MAPPER.readValue(body, ApiErrorBody.class));
+    } catch (IOException e) {
+      return Optional.of(parseUnknownError(response, body, e));
+    }
   }
 
   private static ApiErrorBody parseUnknownError(Response response, String body, IOException err) {

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

@@ -6,6 +6,7 @@
 import com.databricks.sdk.core.error.PrivateLinkValidationError;
 import com.databricks.sdk.core.error.details.ErrorDetails;
 import com.databricks.sdk.core.error.details.ErrorInfo;
+import com.databricks.sdk.core.error.platform.TooManyRequests;
 import com.databricks.sdk.core.http.Request;
 import com.databricks.sdk.core.http.Response;
 import com.databricks.sdk.core.utils.FakeTimer;
@@ -235,8 +236,7 @@ void failAfterTooManyRetries() throws IOException {
             "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);
+    assertInstanceOf(TooManyRequests.class, cause);
   }
 
   @Test
@@ -429,7 +429,7 @@ void verifyNoRetryWhenRetriesDisabled() throws IOException {
                 client.execute(
                     new Request("GET", req.getUri().getPath()), MyEndpointResponse.class));
 
-    assertEquals("TOO_MANY_REQUESTS", exception.getErrorCode());
+    assertInstanceOf(TooManyRequests.class, exception);
     assertEquals(429, exception.getStatusCode());
   }