-
Notifications
You must be signed in to change notification settings - Fork 24.7k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
SQL: forward warning headers to JDBC driver #84499
Changes from 3 commits
990cbcb
056d52c
4c16d66
d3791a3
94d85e0
986ef75
42446b1
4b90b10
321048a
07238bb
a7216d2
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,5 @@ | ||
pr: 84499 | ||
summary: Forward warning headers to JDBC driver | ||
area: SQL | ||
type: enhancement | ||
issues: [] |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -28,4 +28,6 @@ default int columnSize() { | |
int batchSize(); | ||
|
||
void close() throws SQLException; | ||
|
||
List<String> warnings(); | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -33,6 +33,8 @@ | |
import java.io.InputStream; | ||
import java.security.PrivilegedAction; | ||
import java.sql.SQLException; | ||
import java.util.Collections; | ||
import java.util.List; | ||
import java.util.function.Function; | ||
|
||
import static java.util.Collections.emptyList; | ||
|
@@ -82,10 +84,10 @@ public SqlQueryResponse basicQuery(String query, int fetchSize, boolean fieldMul | |
false, | ||
cfg.binaryCommunication() | ||
); | ||
return query(sqlRequest); | ||
return query(sqlRequest).v1(); | ||
} | ||
|
||
public SqlQueryResponse query(SqlQueryRequest sqlRequest) throws SQLException { | ||
public Tuple<SqlQueryResponse, List<String>> query(SqlQueryRequest sqlRequest) throws SQLException { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'm not too fond of working with tuples everywhere I have to pass the warnings but this was the least intrusive way to get the warnings from the response headers. Since it's an internal client it probably doesn't matter too much but suggestions for solving this differently are welcome. |
||
return post(CoreProtocol.SQL_QUERY_REST_ENDPOINT, sqlRequest, Payloads::parseQueryResponse); | ||
} | ||
|
||
|
@@ -98,28 +100,28 @@ public SqlQueryResponse nextPage(String cursor) throws SQLException { | |
new RequestInfo(Mode.CLI), | ||
cfg.binaryCommunication() | ||
); | ||
return post(CoreProtocol.SQL_QUERY_REST_ENDPOINT, sqlRequest, Payloads::parseQueryResponse); | ||
return post(CoreProtocol.SQL_QUERY_REST_ENDPOINT, sqlRequest, Payloads::parseQueryResponse).v1(); | ||
} | ||
|
||
public boolean queryClose(String cursor, Mode mode) throws SQLException { | ||
SqlClearCursorResponse response = post( | ||
Tuple<SqlClearCursorResponse, List<String>> response = post( | ||
CoreProtocol.CLEAR_CURSOR_REST_ENDPOINT, | ||
new SqlClearCursorRequest(cursor, new RequestInfo(mode)), | ||
Payloads::parseClearCursorResponse | ||
); | ||
return response.isSucceeded(); | ||
return response.v1().isSucceeded(); | ||
} | ||
|
||
@SuppressWarnings({ "removal" }) | ||
private <Request extends AbstractSqlRequest, Response> Response post( | ||
private <Request extends AbstractSqlRequest, Response> Tuple<Response, List<String>> post( | ||
String path, | ||
Request request, | ||
CheckedFunction<JsonParser, Response, IOException> responseParser | ||
) throws SQLException { | ||
byte[] requestBytes = toContent(request); | ||
String query = "error_trace"; | ||
Tuple<ContentType, byte[]> response = java.security.AccessController.doPrivileged( | ||
(PrivilegedAction<ResponseOrException<Tuple<ContentType, byte[]>>>) () -> JreHttpUrlConnection.http( | ||
Tuple<Function<String, List<String>>, byte[]> response = java.security.AccessController.doPrivileged( | ||
(PrivilegedAction<ResponseOrException<Tuple<Function<String, List<String>>, byte[]>>>) () -> JreHttpUrlConnection.http( | ||
path, | ||
query, | ||
cfg, | ||
|
@@ -131,7 +133,11 @@ private <Request extends AbstractSqlRequest, Response> Response post( | |
) | ||
) | ||
).getResponseOrThrowException(); | ||
return fromContent(response.v1(), response.v2(), responseParser); | ||
List<String> warnings = response.v1().apply("Warning"); | ||
return new Tuple<>( | ||
fromContent(response.v1(), response.v2(), responseParser), | ||
warnings == null ? Collections.emptyList() : warnings | ||
); | ||
} | ||
|
||
@SuppressWarnings({ "removal" }) | ||
|
@@ -162,8 +168,8 @@ private boolean head(String path, long timeoutInMs) throws SQLException { | |
|
||
@SuppressWarnings({ "removal" }) | ||
private <Response> Response get(String path, CheckedFunction<JsonParser, Response, IOException> responseParser) throws SQLException { | ||
Tuple<ContentType, byte[]> response = java.security.AccessController.doPrivileged( | ||
(PrivilegedAction<ResponseOrException<Tuple<ContentType, byte[]>>>) () -> JreHttpUrlConnection.http( | ||
Tuple<Function<String, List<String>>, byte[]> response = java.security.AccessController.doPrivileged( | ||
(PrivilegedAction<ResponseOrException<Tuple<Function<String, List<String>>, byte[]>>>) () -> JreHttpUrlConnection.http( | ||
path, | ||
"error_trace", | ||
cfg, | ||
|
@@ -184,31 +190,31 @@ private <Request extends AbstractSqlRequest> byte[] toContent(Request request) { | |
} | ||
} | ||
|
||
private Tuple<ContentType, byte[]> readFrom(InputStream inputStream, Function<String, String> headers) { | ||
String contentType = headers.apply("Content-Type"); | ||
ContentType type = ContentFactory.parseMediaType(contentType); | ||
if (type == null) { | ||
throw new IllegalStateException("Unsupported Content-Type: " + contentType); | ||
} | ||
private Tuple<Function<String, List<String>>, byte[]> readFrom(InputStream inputStream, Function<String, List<String>> headers) { | ||
ByteArrayOutputStream out = new ByteArrayOutputStream(); | ||
try { | ||
Streams.copy(inputStream, out); | ||
} catch (Exception ex) { | ||
throw new ClientException("Cannot deserialize response", ex); | ||
} | ||
return new Tuple<>(type, out.toByteArray()); | ||
return new Tuple<>(headers, out.toByteArray()); | ||
|
||
} | ||
|
||
private <Response> Response fromContent( | ||
ContentType contentType, | ||
Function<String, List<String>> headers, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Here I would keep the initial intention of the method (getting the contentType and not the whole list of headers). |
||
byte[] bytesReference, | ||
CheckedFunction<JsonParser, Response, IOException> responseParser | ||
) { | ||
try ( | ||
InputStream stream = new ByteArrayInputStream(bytesReference); | ||
JsonParser parser = ContentFactory.parser(contentType, stream) | ||
) { | ||
List<String> contentTypeHeaders = headers.apply("Content-Type"); | ||
|
||
String contentType = contentTypeHeaders == null || contentTypeHeaders.isEmpty() ? null : contentTypeHeaders.get(0); | ||
ContentType type = ContentFactory.parseMediaType(contentType); | ||
if (type == null) { | ||
throw new IllegalStateException("Unsupported Content-Type: " + contentType); | ||
} | ||
|
||
try (InputStream stream = new ByteArrayInputStream(bytesReference); JsonParser parser = ContentFactory.parser(type, stream)) { | ||
return responseParser.apply(parser); | ||
} catch (Exception ex) { | ||
throw new ClientException("Cannot parse response", ex); | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -20,6 +20,7 @@ | |
import java.net.URI; | ||
import java.net.URISyntaxException; | ||
import java.net.URL; | ||
import java.net.URLConnection; | ||
import java.security.AccessController; | ||
import java.security.PrivilegedAction; | ||
import java.sql.SQLClientInfoException; | ||
|
@@ -31,6 +32,8 @@ | |
import java.sql.SQLSyntaxErrorException; | ||
import java.sql.SQLTimeoutException; | ||
import java.util.Base64; | ||
import java.util.LinkedList; | ||
import java.util.List; | ||
import java.util.function.Function; | ||
import java.util.zip.GZIPInputStream; | ||
|
||
|
@@ -150,15 +153,15 @@ public boolean head() throws ClientException { | |
|
||
public <R> ResponseOrException<R> request( | ||
CheckedConsumer<OutputStream, IOException> doc, | ||
CheckedBiFunction<InputStream, Function<String, String>, R, IOException> parser, | ||
CheckedBiFunction<InputStream, Function<String, List<String>>, R, IOException> parser, | ||
String requestMethod | ||
) throws ClientException { | ||
return request(doc, parser, requestMethod, "application/json"); | ||
} | ||
|
||
public <R> ResponseOrException<R> request( | ||
CheckedConsumer<OutputStream, IOException> doc, | ||
CheckedBiFunction<InputStream, Function<String, String>, R, IOException> parser, | ||
CheckedBiFunction<InputStream, Function<String, List<String>>, R, IOException> parser, | ||
String requestMethod, | ||
String contentTypeHeader | ||
) throws ClientException { | ||
|
@@ -174,7 +177,7 @@ public <R> ResponseOrException<R> request( | |
} | ||
if (shouldParseBody(con.getResponseCode())) { | ||
try (InputStream stream = getStream(con, con.getInputStream())) { | ||
return new ResponseOrException<>(parser.apply(new BufferedInputStream(stream), con::getHeaderField)); | ||
return new ResponseOrException<>(parser.apply(new BufferedInputStream(stream), getHeaderFields(con))); | ||
} | ||
} | ||
return parserError(); | ||
|
@@ -183,6 +186,23 @@ public <R> ResponseOrException<R> request( | |
} | ||
} | ||
|
||
private Function<String, List<String>> getHeaderFields(URLConnection con) { | ||
return header -> { | ||
// HTTP headers are case-insensitive but the map returned by `URLConnection.getHeaderFields` is case-sensitive. | ||
// The linear scan below replicates the linear case-insensitive lookup of `URLConnection.getHeaderField(String)`. | ||
List<String> values = new LinkedList<>(); | ||
int i = 0; | ||
String value = con.getHeaderField(i); | ||
while (value != null) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Just noting that the function would stop iterating on an empty header value? This seems to be legal, though very corner case. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I hope that's never the case but the API does not forbid |
||
if (header.equalsIgnoreCase(con.getHeaderFieldKey(i))) { | ||
values.add(value); | ||
} | ||
value = con.getHeaderField(++i); | ||
} | ||
return values; | ||
}; | ||
} | ||
|
||
private boolean shouldParseBody(int responseCode) { | ||
return responseCode == 200 || responseCode == 201 || responseCode == 202; | ||
} | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please, move this line at the top of the class.