Chunk streaming request bodies only

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

@@ -12,10 +12,9 @@
 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;
@@ -187,12 +186,10 @@ private <I> Request prepareBaseRequest(String method, String path, I in)
       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);
+      return new Request(method, path, body);
     } else {
-      String debugBody = serialize(in);
-      InputStream body = new ByteArrayInputStream(debugBody.getBytes(StandardCharsets.UTF_8));
-      return new Request(method, path, body, debugBody);
+      String body = serialize(in);
+      return new Request(method, path, body);
     }
   }
 
@@ -303,11 +300,15 @@ private String makeLogRecord(Request in, Response out) {
       in.getHeaders()
           .forEach((header, value) -> sb.append(String.format("\n * %s: %s", header, value)));
     }
-    String requestBody = in.getDebugBody();
-    if (requestBody != null && !requestBody.isEmpty()) {
-      for (String line : bodyLogger.redactedDump(requestBody).split("\n")) {
-        sb.append("\n> ");
-        sb.append(line);
+    if (in.isBodyStreaming()) {
+      sb.append("\n> (streamed body)");
+    } else {
+      String requestBody = in.getBodyString();
+      if (requestBody != null && !requestBody.isEmpty()) {
+        for (String line : bodyLogger.redactedDump(requestBody).split("\n")) {
+          sb.append("\n> ");
+          sb.append(line);
+        }
       }
     }
     sb.append("\n< ");

databricks-sdk-java/src/main/java/com/databricks/sdk/core/commons/CommonsHttpClient.java

@@ -14,18 +14,24 @@
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
+
+import com.databricks.sdk.mixin.ClustersExt;
 import org.apache.commons.io.IOUtils;
 import org.apache.http.HttpEntity;
 import org.apache.http.NameValuePair;
 import org.apache.http.StatusLine;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.*;
 import org.apache.http.entity.InputStreamEntity;
+import org.apache.http.entity.StringEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class CommonsHttpClient implements HttpClient {
+  private static final Logger LOG = LoggerFactory.getLogger(ClustersExt.class);
   private final PoolingHttpClientConnectionManager connectionManager =
       new PoolingHttpClientConnectionManager();
   private final CloseableHttpClient hc;
@@ -117,18 +123,24 @@ private HttpUriRequest transformRequest(Request in) {
       case Request.DELETE:
         return new HttpDelete(in.getUri());
       case Request.POST:
-        return withEntity(new HttpPost(in.getUri()), in.getBody());
+        return withEntity(new HttpPost(in.getUri()), in);
       case Request.PUT:
-        return withEntity(new HttpPut(in.getUri()), in.getBody());
+        return withEntity(new HttpPut(in.getUri()), in);
       case Request.PATCH:
-        return withEntity(new HttpPatch(in.getUri()), in.getBody());
+        return withEntity(new HttpPatch(in.getUri()), in);
       default:
         throw new IllegalArgumentException("Unknown method: " + in.getMethod());
     }
   }
 
-  private HttpRequestBase withEntity(HttpEntityEnclosingRequestBase request, InputStream body) {
-    request.setEntity(new InputStreamEntity(body));
+  private HttpRequestBase withEntity(HttpEntityEnclosingRequestBase request, Request in) {
+    if (in.isBodyString()) {
+      request.setEntity(new StringEntity(in.getBodyString(), StandardCharsets.UTF_8));
+    } else if (in.isBodyStreaming()) {
+      request.setEntity(new InputStreamEntity(in.getBodyStream()));
+    } else {
+      LOG.warn("withEntity called with a request with no body, so no request entity will be set. URI: {}", in.getUri());
+    }
     return request;
   }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/http/Request.java

@@ -1,7 +1,7 @@
 package com.databricks.sdk.core.http;
 
 import com.databricks.sdk.core.DatabricksException;
-import java.io.ByteArrayInputStream;
+
 import java.io.InputStream;
 import java.io.UnsupportedEncodingException;
 import java.net.URI;
@@ -19,26 +19,49 @@ public class Request {
   private String url;
   private final Map<String, String> headers = new HashMap<>();
   private final Map<String, List<String>> query = new TreeMap<>();
-  private final InputStream body;
-  private final String debugBody;
+  /**
+   * The body of the request for requests with streaming bodies. At most one of {@link #bodyStream}
+   * and {@link #bodyString} can be non-null.
+   */
+  private final InputStream bodyStream;
+  /**
+   * The body of the request for requests with string bodies. At most one of {@link #bodyStream}
+   * and {@link #bodyString} can be non-null.
+   */
+  private final String bodyString;
+  /**
+   * Whether the body of the request is a streaming body. At most one of {@link #isBodyStreaming} and
+   * {@link #isBodyString} can be true.
+   */
+  private final boolean isBodyStreaming;
+  /**
+   * Whether the body of the request is a string body. At most one of {@link #isBodyStreaming} and
+   * {@link #isBodyString} can be true.
+   */
+  private final boolean isBodyString;
 
   public Request(String method, String url) {
-    this(method, url, (String) null);
+    this(method, url, null, null);
   }
 
-  public Request(String method, String url, String body) {
-    this.method = method;
-    this.url = url;
-    this.body =
-        body != null ? new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8)) : null;
-    this.debugBody = body;
+  public Request(String method, String url, String bodyString) {
+    this(method, url, null, bodyString);
+  }
+
+  public Request(String method, String url, InputStream bodyStream) {
+    this(method, url, bodyStream, null);
   }
 
-  public Request(String method, String url, InputStream body, String debugBody) {
+  private Request(String method, String url, InputStream bodyStream, String bodyString) {
+    if (bodyStream != null && bodyString != null) {
+      throw new IllegalArgumentException("At most one of bodyStream and bodyString can be non-null");
+    }
     this.method = method;
     this.url = url;
-    this.body = body;
-    this.debugBody = debugBody;
+    this.bodyStream = bodyStream;
+    this.bodyString = bodyString;
+    this.isBodyStreaming = bodyStream != null;
+    this.isBodyString = bodyString != null;
   }
 
   public Request withHeaders(Map<String, String> headers) {
@@ -135,12 +158,20 @@ public Map<String, List<String>> getQuery() {
     return query;
   }
 
-  public InputStream getBody() {
-    return body;
+  public InputStream getBodyStream() {
+    return bodyStream;
+  }
+
+  public String getBodyString() {
+    return bodyString;
+  }
+
+  public boolean isBodyStreaming() {
+    return isBodyStreaming;
   }
 
-  public String getDebugBody() {
-    return debugBody;
+  public boolean isBodyString() {
+    return isBodyString;
   }
 
   @Override
@@ -151,15 +182,15 @@ public boolean equals(Object o) {
     return method.equals(request.method)
         && url.equals(request.url)
         && Objects.equals(query, request.query)
-        && Objects.equals(body, request.body);
+        && Objects.equals(bodyStream, request.bodyStream);
   }
 
   @Override
   public int hashCode() {
     // Note: this is not safe for production, as debugBody will be the same for different requests
     // using InputStream.
     // It is currently only used in tests.
-    return Objects.hash(method, url, query, debugBody);
+    return Objects.hash(method, url, query, bodyString);
   }
 
   @Override