databricks · mgyucht · Aug 23, 2023 · Aug 14, 2023 · Aug 14, 2023 · Aug 15, 2023
diff --git a/.codegen/api.java.tmpl b/.codegen/api.java.tmpl
@@ -2,6 +2,7 @@
 package com.databricks.sdk.service.{{.Package.Name}};
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.Collection;
 import java.util.Map;
 import java.time.Duration;

diff --git a/.codegen/impl.java.tmpl b/.codegen/impl.java.tmpl
@@ -2,6 +2,7 @@
 package com.databricks.sdk.service.{{.Package.Name}};
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.Collection;
 import java.util.Map;
 import java.util.HashMap;
@@ -27,23 +28,33 @@ class {{.PascalName}}Impl implements {{.PascalName}}Service {
     String path = {{if .PathParts -}}
       String.format("{{range  .PathParts}}{{.Prefix}}{{if or .Field .IsAccountId}}%s{{end}}{{ end }}"{{ range .PathParts }}{{if .Field}}, request.get{{.Field.PascalName}}(){{ else if .IsAccountId }}, apiClient.configuredAccountID(){{end}}{{ end }})
     {{- else}}"{{.Path}}"{{end}};
-    {{ template "headers" . }}
-    {{if .Response -}}
-      {{- if .Response.ArrayValue -}} return apiClient.getCollection(path, null, {{template "type" .Response.ArrayValue}}.class, headers);
-      {{- else if .Response.MapValue -}} return apiClient.getStringMap(path, request, headers);
-      {{- else -}} return apiClient.{{.Verb}}(path{{if .Request}}, request{{end}}{{if .Response -}}, {{if .Response.ArrayValue }}Collection
-                        {{- else if .Response.MapValue }}Map
-                        {{- else}}{{template "type" .Response}}
-                        {{- end -}}{{else}}, Void{{end}}.class, headers);{{end}}
-    {{- else -}}apiClient.{{.Verb}}(path{{if .Request}}, request{{end}}{{if .Response -}}, {{if .Response.ArrayValue }}Collection
-      {{- else if .Response.MapValue }}Map
-      {{- else}}{{template "type" .Response}}
-      {{- end -}}{{else}}, Void{{end}}.class, headers);
-    {{end}}
+    {{ template "headers" . -}}
+    {{ if not .Response -}}
+    {{ template "api-call" . }}
+    {{- else if .Response.ArrayValue -}} return apiClient.getCollection(path, null, {{template "type" .Response.ArrayValue}}.class, headers);
+    {{- else if .Response.MapValue -}} return apiClient.getStringMap(path, {{ template "request-param" .}}, headers);
+    {{- else if .IsResponseByteStream -}}
+    InputStream response = {{ template "api-call" . }}
+    return new {{ .Response.PascalName }}().set{{.ResponseBodyField.PascalName}}(response);
+    {{- else }}return {{ template "api-call" . }}
+    {{- end}}
   }
   {{end}}
 }
 
+{{ define "api-call" }}
+apiClient.{{.Verb}}(path
+  {{- if .Request}}, {{ template "request-param" .}}{{end}}
+  , {{ if not .Response -}}Void
+  {{- else if .IsResponseByteStream}}InputStream
+  {{- else}}{{template "type" .Response}}{{- end -}}.class
+  , headers);
+{{- end }}
+
+{{ define "request-param" -}}
+  request{{ if .RequestBodyField }}.get{{.RequestBodyField.PascalName}}(){{end}}
+{{- end }}
+
 {{ define "headers" -}}
     Map<String, String> headers = new HashMap<>();
     {{- range $key, $value := .FixedRequestHeaders }}

diff --git a/.codegen/interface.java.tmpl b/.codegen/interface.java.tmpl
@@ -3,6 +3,7 @@ package com.databricks.sdk.service.{{.Package.Name}};
 
 import java.util.Collection;
 import java.util.Map;
+import java.io.InputStream;
 
 import com.databricks.sdk.support.Generated;
 
@@ -33,6 +34,7 @@ public interface {{.PascalName}}Service {
   {{- else if .IsInt}}long
   {{- else if .ArrayValue }}Collection<{{template "type" .ArrayValue}}>
   {{- else if .MapValue }}Map<String,{{template "type" .MapValue}}>
+  {{- else if .IsByteStream}}InputStream
   {{- else if .IsObject }}{{.PascalName}}
   {{- else if .IsExternal }}com.databricks.sdk.service.{{.Package.Name}}.{{.PascalName}}
   {{- else if .Enum }}{{.PascalName}}

diff --git a/.codegen/model.java.tmpl b/.codegen/model.java.tmpl
@@ -3,6 +3,7 @@
 package com.databricks.sdk.service.{{.Package.Name}};
 
 import java.util.Map;
+import java.io.InputStream;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import java.util.Collection;
@@ -87,6 +88,7 @@ public enum {{.PascalName}}{
   {{- else if .IsInt}}Long
   {{- else if .ArrayValue }}Collection<{{template "type" .ArrayValue}}>
   {{- else if .MapValue }}Map<String,{{template "type" .MapValue}}>
+  {{- else if .IsByteStream}}InputStream
   {{- else if .IsObject }}{{.PascalName}}
   {{- else if .IsExternal }}com.databricks.sdk.service.{{.Package.Name}}.{{.PascalName}}
   {{- else if .Enum }}{{.PascalName}}

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
@@ -12,7 +12,10 @@
 import com.fasterxml.jackson.databind.JavaType;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.SerializationFeature;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -182,9 +185,14 @@ private <I> Request prepareBaseRequest(String method, String path, I in)
       throws JsonProcessingException {
     if (in == null || !hasBody(method)) {
       return new Request(method, path);
+    } else if (InputStream.class.isAssignableFrom(in.getClass())) {
+      InputStream body = (InputStream) in;
+      String debugBody = "<InputStream>";
+      return new Request(method, path, body, debugBody);
     } else {
-      String body = serialize(in);
-      return new Request(method, path, body);
+      String debugBody = serialize(in);
+      InputStream body = new ByteArrayInputStream(debugBody.getBytes(StandardCharsets.UTF_8));
+      return new Request(method, path, body, debugBody);
     }
   }
 
@@ -295,7 +303,7 @@ private String makeLogRecord(Request in, Response out) {
       in.getHeaders()
           .forEach((header, value) -> sb.append(String.format("\n * %s: %s", header, value)));
     }
-    String requestBody = in.getBody();
+    String requestBody = in.getDebugBody();
     if (requestBody != null && !requestBody.isEmpty()) {
       for (String line : bodyLogger.redactedDump(requestBody).split("\n")) {
         sb.append("\n> ");
@@ -304,18 +312,24 @@ private String makeLogRecord(Request in, Response out) {
     }
     sb.append("\n< ");
     sb.append(out.toString());
-    for (String line : bodyLogger.redactedDump(out.getBody()).split("\n")) {
+    for (String line : bodyLogger.redactedDump(out.getDebugBody()).split("\n")) {
       sb.append("\n< ");
       sb.append(line);
     }
     return sb.toString();
   }
 
-  public <T> T deserialize(String body, Class<T> target) throws IOException {
+  public <T> T deserialize(InputStream body, Class<T> target) throws IOException {
+    if (target == InputStream.class) {
+      return (T) body;
+    }
     return mapper.readValue(body, target);
   }
 
-  public <T> T deserialize(String body, JavaType target) throws IOException {
+  public <T> T deserialize(InputStream body, JavaType target) throws IOException {
+    if (target == mapper.constructType(InputStream.class)) {
+      return (T) body;
+    }
     return mapper.readValue(body, target);
   }
 

diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/BodyLogger.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/BodyLogger.java
@@ -48,8 +48,8 @@ public String redactedDump(String body) {
       Object result = recursiveMarshal(rootNode, maxBytes);
       return mapper.writerWithDefaultPrettyPrinter().writeValueAsString(result);
     } catch (JsonProcessingException e) {
-      // Unable to unmarshal means the body isn't JSON (or something else)
-      return String.format("[unable to marshal: %s]", e.getMessage());
+      // Unable to unmarshal means the body isn't JSON
+      return onlyNBytes(body, maxBytes);
     }
   }
 

diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/commons/CommonsHttpClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/commons/CommonsHttpClient.java
@@ -1,12 +1,15 @@
 package com.databricks.sdk.core.commons;
 
+import static org.apache.http.entity.ContentType.APPLICATION_JSON;
+
+import com.databricks.sdk.core.DatabricksException;
 import com.databricks.sdk.core.http.HttpClient;
 import com.databricks.sdk.core.http.Request;
 import com.databricks.sdk.core.http.Response;
+import com.databricks.sdk.core.utils.CustomCloseInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.UnsupportedEncodingException;
-import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -17,7 +20,7 @@
 import org.apache.http.StatusLine;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.*;
-import org.apache.http.entity.StringEntity;
+import org.apache.http.entity.InputStreamEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
@@ -53,22 +56,56 @@ private CloseableHttpClient makeClosableHttpClient() {
   public Response execute(Request in) throws IOException {
     HttpUriRequest request = transformRequest(in);
     in.getHeaders().forEach(request::setHeader);
-    try (CloseableHttpResponse response = hc.execute(request)) {
-      HttpEntity entity = response.getEntity();
-      StatusLine statusLine = response.getStatusLine();
-      Map<String, List<String>> hs =
-          Arrays.stream(response.getAllHeaders())
-              .collect(
-                  Collectors.groupingBy(
-                      NameValuePair::getName,
-                      Collectors.mapping(NameValuePair::getValue, Collectors.toList())));
-      String body = null;
-      if (entity != null) {
-        try (InputStream inputStream = entity.getContent()) {
-          body = IOUtils.toString(inputStream, Charset.defaultCharset());
-        }
-      }
+    CloseableHttpResponse response = hc.execute(request);
+    return computeResponse(in, response);
+  }
+
+  private Response computeResponse(Request in, CloseableHttpResponse response) throws IOException {
+    HttpEntity entity = response.getEntity();
+    StatusLine statusLine = response.getStatusLine();
+    Map<String, List<String>> hs =
+        Arrays.stream(response.getAllHeaders())
+            .collect(
+                Collectors.groupingBy(
+                    NameValuePair::getName,
+                    Collectors.mapping(NameValuePair::getValue, Collectors.toList())));
+    if (entity == null) {
+      response.close();
+      return new Response(in, statusLine.getStatusCode(), statusLine.getReasonPhrase(), hs);
+    }
+
+    // The Databricks SDK is currently designed to treat all non-application/json responses as
+    // InputStreams, leaving the caller to decide how to read and parse the response. The caller
+    // is responsible for closing the InputStream to release the HTTP Connection.
+    //
+    // The client only streams responses when the caller has explicitly requested a non-JSON
+    // response and the server has responded with a non-JSON Content-Type. The Databricks API
+    // error response is either JSON or HTML and is safe to read fully into memory.
+    boolean streamResponse =
+        in.getHeaders().containsKey("Accept")
+            && !APPLICATION_JSON.getMimeType().equals(in.getHeaders().get("Accept"))
+            && hs.containsKey("Content-Type")
+            && !APPLICATION_JSON.getMimeType().equals(hs.get("Content-Type").get(0));
+    if (streamResponse) {
+      CustomCloseInputStream inputStream =
+          new CustomCloseInputStream(
+              entity.getContent(),
+              () -> {
+                try {
+                  response.close();
+                } catch (Exception e) {
+                  throw new DatabricksException("Unable to close connection", e);
+                }
+              });
+      return new Response(
+          in, statusLine.getStatusCode(), statusLine.getReasonPhrase(), hs, inputStream);
+    }
+
+    try (InputStream inputStream = entity.getContent()) {
+      String body = IOUtils.toString(inputStream, StandardCharsets.UTF_8);
       return new Response(in, statusLine.getStatusCode(), statusLine.getReasonPhrase(), hs, body);
+    } finally {
+      response.close();
     }
   }
 
@@ -89,12 +126,8 @@ private HttpUriRequest transformRequest(Request in) {
     }
   }
 
-  private HttpRequestBase withEntity(HttpEntityEnclosingRequestBase request, String body) {
-    try {
-      request.setEntity(new StringEntity(body));
-      return request;
-    } catch (UnsupportedEncodingException e) {
-      throw new IllegalArgumentException(e);
-    }
+  private HttpRequestBase withEntity(HttpEntityEnclosingRequestBase request, InputStream body) {
+    request.setEntity(new InputStreamEntity(body));
+    return request;
   }
 }
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/error/ApiErrors.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/error/ApiErrors.java
@@ -1,11 +1,14 @@
 package com.databricks.sdk.core.error;
 
 import com.databricks.sdk.core.DatabricksError;
+import com.databricks.sdk.core.DatabricksException;
 import com.databricks.sdk.core.http.Response;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import java.io.IOException;
+import java.io.*;
+import java.nio.charset.StandardCharsets;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import org.apache.commons.io.IOUtils;
 
 /** Helper methods for inspecting the response and errors thrown during API requests. */
 public class ApiErrors {
@@ -55,13 +58,20 @@ private static DatabricksError readErrorFromResponse(Response response) {
    */
   private static ApiErrorBody parseApiError(Response response) {
     try {
-      return MAPPER.readValue(response.getBody(), ApiErrorBody.class);
+      // Read the body now, so we can try to parse as JSON and then fallback to old error handling
+      // logic.
+      String body = IOUtils.toString(response.getBody(), StandardCharsets.UTF_8);
+      try {
+        return MAPPER.readValue(body, ApiErrorBody.class);
+      } catch (IOException e) {
+        return parseUnknownError(response, body, e);
+      }
     } catch (IOException e) {
-      return parseUnknownError(response, e);
+      throw new DatabricksException("Unable to read response body", e);
     }
   }
 
-  private static ApiErrorBody parseUnknownError(Response response, IOException err) {
+  private static ApiErrorBody parseUnknownError(Response response, String body, IOException err) {
     ApiErrorBody errorBody = new ApiErrorBody();
     String[] statusParts = response.getStatus().split(" ", 2);
     if (statusParts.length < 2) {
@@ -71,14 +81,13 @@ private static ApiErrorBody parseUnknownError(Response response, IOException err
       errorBody.setErrorCode(errorCode.replaceAll(" ", "_").toUpperCase());
     }
 
-    Matcher messageMatcher = HTML_ERROR_REGEX.matcher(response.getBody());
+    Matcher messageMatcher = HTML_ERROR_REGEX.matcher(body);
     if (messageMatcher.find()) {
       errorBody.setMessage(messageMatcher.group(1).replaceAll("^[ .]+|[ .]+$", ""));
     } else {
       errorBody.setMessage(
           String.format(
-              "Response from server (%s) %s: %s",
-              response.getStatus(), response.getBody(), err.getMessage()));
+              "Response from server (%s) %s: %s", response.getStatus(), body, err.getMessage()));
     }
     return errorBody;
   }