Skip to content

Commit 6b92a5a

Browse files
authored
Chunk streaming request bodies only (#157)
## Changes #138 changed the CommonsHttpClient implementation to use `InputStreamEntity` for all request bodies sent using the Java SDK. As a result, requests are sent with a `Transfer-Encoding: chunked` header, and the request body is chunked accordingly. However, the Databricks REST API only tolerates chunked requests for certain APIs; for others, it ignores the request body if Content-Length is not specified. Auth falls under that category, which caused #154. To fix this, Request will keep track of string entities and input stream entities separately. Previously, Request had separate fields for the body and debugBody, but they were both always set, so it wasn't possible to distinguish between these two cases. Now, HTTP clients can have different behavior based on whether the request body can be fully materialized as a string or is lazily read as with input streams. As part of this, I have removed `SimpleHttpServer`. This was used before we were sure whether it was possible to use the HttpServer built into the JDK. Now that we are confident that we can (see usage in ExternalBrowserCredentialsProvider.java), I've replaced it with HttpServer and a custom HttpHandler in FixtureServer (the only place where it is used now). One other small change: I've updated the logging configuration and dependency for the cli auth demo app. This ensures that users can see debug output (which I used when debugging this). Closes #154. ## Tests I've refactored FixtureServer to support more advanced validation on the HTTP requests sent by clients for testing purposes. Now, users can assert that a request has a specific method, path, headers, body, and check for headers that should not be present. One downside of this change is that users need to call the `.with()` method once per API call, rather than one time, at the start of each test. Unit tests cover both cases (request body should be chunked when specified as a String; request body should not be chunked when specified as an InputStream).
1 parent 00052d4 commit 6b92a5a

File tree

11 files changed

+346
-238
lines changed

11 files changed

+346
-238
lines changed

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

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,8 @@
1212
import com.fasterxml.jackson.databind.JavaType;
1313
import com.fasterxml.jackson.databind.ObjectMapper;
1414
import com.fasterxml.jackson.databind.SerializationFeature;
15-
import java.io.ByteArrayInputStream;
1615
import java.io.IOException;
1716
import java.io.InputStream;
18-
import java.nio.charset.StandardCharsets;
1917
import java.util.*;
2018
import org.slf4j.Logger;
2119
import org.slf4j.LoggerFactory;
@@ -187,12 +185,10 @@ private <I> Request prepareBaseRequest(String method, String path, I in)
187185
return new Request(method, path);
188186
} else if (InputStream.class.isAssignableFrom(in.getClass())) {
189187
InputStream body = (InputStream) in;
190-
String debugBody = "<InputStream>";
191-
return new Request(method, path, body, debugBody);
188+
return new Request(method, path, body);
192189
} else {
193-
String debugBody = serialize(in);
194-
InputStream body = new ByteArrayInputStream(debugBody.getBytes(StandardCharsets.UTF_8));
195-
return new Request(method, path, body, debugBody);
190+
String body = serialize(in);
191+
return new Request(method, path, body);
196192
}
197193
}
198194

@@ -303,11 +299,15 @@ private String makeLogRecord(Request in, Response out) {
303299
in.getHeaders()
304300
.forEach((header, value) -> sb.append(String.format("\n * %s: %s", header, value)));
305301
}
306-
String requestBody = in.getDebugBody();
307-
if (requestBody != null && !requestBody.isEmpty()) {
308-
for (String line : bodyLogger.redactedDump(requestBody).split("\n")) {
309-
sb.append("\n> ");
310-
sb.append(line);
302+
if (in.isBodyStreaming()) {
303+
sb.append("\n> (streamed body)");
304+
} else {
305+
String requestBody = in.getBodyString();
306+
if (requestBody != null && !requestBody.isEmpty()) {
307+
for (String line : bodyLogger.redactedDump(requestBody).split("\n")) {
308+
sb.append("\n> ");
309+
sb.append(line);
310+
}
311311
}
312312
}
313313
sb.append("\n< ");

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

Lines changed: 0 additions & 108 deletions
This file was deleted.

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

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,15 @@
2121
import org.apache.http.client.config.RequestConfig;
2222
import org.apache.http.client.methods.*;
2323
import org.apache.http.entity.InputStreamEntity;
24+
import org.apache.http.entity.StringEntity;
2425
import org.apache.http.impl.client.CloseableHttpClient;
2526
import org.apache.http.impl.client.HttpClientBuilder;
2627
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
28+
import org.slf4j.Logger;
29+
import org.slf4j.LoggerFactory;
2730

2831
public class CommonsHttpClient implements HttpClient {
32+
private static final Logger LOG = LoggerFactory.getLogger(CommonsHttpClient.class);
2933
private final PoolingHttpClientConnectionManager connectionManager =
3034
new PoolingHttpClientConnectionManager();
3135
private final CloseableHttpClient hc;
@@ -117,18 +121,26 @@ private HttpUriRequest transformRequest(Request in) {
117121
case Request.DELETE:
118122
return new HttpDelete(in.getUri());
119123
case Request.POST:
120-
return withEntity(new HttpPost(in.getUri()), in.getBody());
124+
return withEntity(new HttpPost(in.getUri()), in);
121125
case Request.PUT:
122-
return withEntity(new HttpPut(in.getUri()), in.getBody());
126+
return withEntity(new HttpPut(in.getUri()), in);
123127
case Request.PATCH:
124-
return withEntity(new HttpPatch(in.getUri()), in.getBody());
128+
return withEntity(new HttpPatch(in.getUri()), in);
125129
default:
126130
throw new IllegalArgumentException("Unknown method: " + in.getMethod());
127131
}
128132
}
129133

130-
private HttpRequestBase withEntity(HttpEntityEnclosingRequestBase request, InputStream body) {
131-
request.setEntity(new InputStreamEntity(body));
134+
private HttpRequestBase withEntity(HttpEntityEnclosingRequestBase request, Request in) {
135+
if (in.isBodyString()) {
136+
request.setEntity(new StringEntity(in.getBodyString(), StandardCharsets.UTF_8));
137+
} else if (in.isBodyStreaming()) {
138+
request.setEntity(new InputStreamEntity(in.getBodyStream()));
139+
} else {
140+
LOG.warn(
141+
"withEntity called with a request with no body, so no request entity will be set. URI: {}",
142+
in.getUri());
143+
}
132144
return request;
133145
}
134146
}

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

Lines changed: 51 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package com.databricks.sdk.core.http;
22

33
import com.databricks.sdk.core.DatabricksException;
4-
import java.io.ByteArrayInputStream;
54
import java.io.InputStream;
65
import java.io.UnsupportedEncodingException;
76
import java.net.URI;
@@ -19,26 +18,50 @@ public class Request {
1918
private String url;
2019
private final Map<String, String> headers = new HashMap<>();
2120
private final Map<String, List<String>> query = new TreeMap<>();
22-
private final InputStream body;
23-
private final String debugBody;
21+
/**
22+
* The body of the request for requests with streaming bodies. At most one of {@link #bodyStream}
23+
* and {@link #bodyString} can be non-null.
24+
*/
25+
private final InputStream bodyStream;
26+
/**
27+
* The body of the request for requests with string bodies. At most one of {@link #bodyStream} and
28+
* {@link #bodyString} can be non-null.
29+
*/
30+
private final String bodyString;
31+
/**
32+
* Whether the body of the request is a streaming body. At most one of {@link #isBodyStreaming}
33+
* and {@link #isBodyString} can be true.
34+
*/
35+
private final boolean isBodyStreaming;
36+
/**
37+
* Whether the body of the request is a string body. At most one of {@link #isBodyStreaming} and
38+
* {@link #isBodyString} can be true.
39+
*/
40+
private final boolean isBodyString;
2441

2542
public Request(String method, String url) {
26-
this(method, url, (String) null);
43+
this(method, url, null, null);
2744
}
2845

29-
public Request(String method, String url, String body) {
30-
this.method = method;
31-
this.url = url;
32-
this.body =
33-
body != null ? new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8)) : null;
34-
this.debugBody = body;
46+
public Request(String method, String url, String bodyString) {
47+
this(method, url, null, bodyString);
3548
}
3649

37-
public Request(String method, String url, InputStream body, String debugBody) {
50+
public Request(String method, String url, InputStream bodyStream) {
51+
this(method, url, bodyStream, null);
52+
}
53+
54+
private Request(String method, String url, InputStream bodyStream, String bodyString) {
55+
if (bodyStream != null && bodyString != null) {
56+
throw new IllegalArgumentException(
57+
"At most one of bodyStream and bodyString can be non-null");
58+
}
3859
this.method = method;
3960
this.url = url;
40-
this.body = body;
41-
this.debugBody = debugBody;
61+
this.bodyStream = bodyStream;
62+
this.bodyString = bodyString;
63+
this.isBodyStreaming = bodyStream != null;
64+
this.isBodyString = bodyString != null;
4265
}
4366

4467
public Request withHeaders(Map<String, String> headers) {
@@ -135,12 +158,20 @@ public Map<String, List<String>> getQuery() {
135158
return query;
136159
}
137160

138-
public InputStream getBody() {
139-
return body;
161+
public InputStream getBodyStream() {
162+
return bodyStream;
163+
}
164+
165+
public String getBodyString() {
166+
return bodyString;
167+
}
168+
169+
public boolean isBodyStreaming() {
170+
return isBodyStreaming;
140171
}
141172

142-
public String getDebugBody() {
143-
return debugBody;
173+
public boolean isBodyString() {
174+
return isBodyString;
144175
}
145176

146177
@Override
@@ -151,15 +182,15 @@ public boolean equals(Object o) {
151182
return method.equals(request.method)
152183
&& url.equals(request.url)
153184
&& Objects.equals(query, request.query)
154-
&& Objects.equals(body, request.body);
185+
&& Objects.equals(bodyStream, request.bodyStream);
155186
}
156187

157188
@Override
158189
public int hashCode() {
159-
// Note: this is not safe for production, as debugBody will be the same for different requests
190+
// Note: this is not safe for production, as bodyString will be null for different requests
160191
// using InputStream.
161192
// It is currently only used in tests.
162-
return Objects.hash(method, url, query, debugBody);
193+
return Objects.hash(method, url, query, bodyString);
163194
}
164195

165196
@Override

0 commit comments

Comments
 (0)