Support inputstream for request and response bodies (#138)

## Changes
This PR adds support for non-application/json requests and responses.
This PR is divided into two parts:

**Support setting headers per request.** The main methods of the
ApiClient interface are expanded to accept one new
parameter,`Map<String, String> headers`, and all callers are expected to
pass a map of headers to be included in the request. As the allowed
content types for requests and responses are known from the OpenAPI
specification, impls must construct the request header map in code
generation and pass it to the ApiClient. The default "Content-Type:
application/json" header is removed, as each request should specify its
own Content-Type and Accept headers. (This is merged in
#135.)

**Add support for streaming request and response bodies to support
non-application/json requests/responses.** Today, serialization of
requests is done in the ApiClient. This implies that ApiClient needs to
be able to serialize a request purely based on the parameters provided
to it via headers, the request body, etc. In this proposal, the
serialization is coupled to the request type: InputStream indicates that
the client has already serialized the request, so it can be provided
directly to the underlying HttpClient library, and any other type will
be passed through Jackson's ObjectMapper for serialization to JSON.
Internally, request and response bodies are modeled as InputStreams
rather than byte arrays to support streaming requests and responses,
though for non-streaming requests and responses, a string value for the
body is captured in the `debugBody` field of Request and Response,
respectively.

One important caveat about streaming responses is that callers are
required to close the streams when they are done using them. This
releases the underlying HTTP connection back to the connection pool for
use in subsequent requests. Otherwise, customers could encounter
deadlocks waiting for connections. For this, we recommend that users use
the try-with-resources model as demonstrated in integration test
FilesIT.java.

## Tests
- [x] Integration test for Files IT exercises uploading and downloading
files using the streaming pathways.
- [x] Existing integration tests cover the pre-existing non-streaming
pathways.

Author: mgyucht

.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;

.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 }}

.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}}

.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}}

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);
   }
 

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);
     }
   }
 

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;
   }
 }

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;
   }