Skip to content

Commit a9bbebe

Browse files
bpinteabpintea
authored
SQL: Add a media type parser for SQL requests (#74116)
This adds a (branch-specific) media type parser as a near-drop-in replacement for the (master-only) per-REST-endpoint media types parser (#64406).
1 parent ad0f96a commit a9bbebe

File tree

2 files changed

+251
-0
lines changed

2 files changed

+251
-0
lines changed

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/SqlMediaTypeParser.java‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.sql.plugin;
9+
10+
import org.elasticsearch.common.Strings;
11+
import org.elasticsearch.common.xcontent.XContentType;
12+
import org.elasticsearch.core.Nullable;
13+
import org.elasticsearch.rest.RestRequest;
14+
import org.elasticsearch.xpack.sql.action.SqlQueryRequest;
15+
import org.elasticsearch.xpack.sql.proto.Mode;
16+
17+
import java.util.List;
18+
import java.util.Locale;
19+
import java.util.Map;
20+
21+
import static org.elasticsearch.xpack.sql.proto.Protocol.URL_PARAM_FORMAT;
22+
23+
public class SqlMediaTypeParser {
24+
25+
static class SqlMediaType {
26+
private final XContentType xContentType;
27+
private final TextFormat textFormat;
28+
private final boolean isTextFormat;
29+
30+
private SqlMediaType(XContentType xContentType, TextFormat textFormat) {
31+
this.xContentType = xContentType;
32+
this.textFormat = textFormat;
33+
isTextFormat = textFormat != null;
34+
}
35+
36+
private static SqlMediaType xContentType(XContentType xContentType) {
37+
return xContentType != null ? new SqlMediaType(xContentType, null) : null;
38+
}
39+
40+
private static SqlMediaType textFormat(TextFormat textFormat) {
41+
return textFormat != null ? new SqlMediaType(null, textFormat) : null;
42+
}
43+
44+
boolean isTextFormat() {
45+
return isTextFormat;
46+
}
47+
48+
XContentType xContentType() {
49+
return xContentType;
50+
}
51+
52+
TextFormat textFormat() {
53+
return textFormat;
54+
}
55+
56+
private static SqlMediaType fromMediaTypeOrFormat(String mediaType) {
57+
XContentType xContentType = XContentType.fromMediaTypeOrFormat(mediaType);
58+
return xContentType != null ? xContentType(xContentType) : textFormat(TextFormat.fromMediaTypeOrFormat(mediaType));
59+
}
60+
}
61+
/*
62+
* Since we support {@link TextFormat} <strong>and</strong>
63+
* {@link XContent} outputs we can't use {@link RestToXContentListener}
64+
* like everything else. We want to stick as closely as possible to
65+
* Elasticsearch's defaults though, while still layering in ways to
66+
* control the output more easily.
67+
*
68+
* First we find the string that the user used to specify the response
69+
* format. If there is a {@code format} parameter we use that. If there
70+
* isn't but there is a {@code Accept} header then we use that. If there
71+
* isn't then we use the {@code Content-Type} header which is required.
72+
*/
73+
public static SqlMediaType getResponseMediaType(RestRequest request, SqlQueryRequest sqlRequest) {
74+
if (Mode.isDedicatedClient(sqlRequest.requestInfo().mode())
75+
&& (sqlRequest.binaryCommunication() == null || sqlRequest.binaryCommunication())) {
76+
// enforce CBOR response for drivers and CLI (unless instructed differently through the config param)
77+
return SqlMediaType.xContentType(XContentType.CBOR);
78+
} else if (request.hasParam(URL_PARAM_FORMAT)) {
79+
return validateColumnarRequest(sqlRequest.columnar(), mediaTypeFromParams(request), request);
80+
}
81+
82+
return mediaTypeFromHeaders(request);
83+
}
84+
85+
public static SqlMediaType getResponseMediaType(RestRequest request) {
86+
return request.hasParam(URL_PARAM_FORMAT)
87+
? checkNonNullMediaType(mediaTypeFromParams(request), request)
88+
: mediaTypeFromHeaders(request);
89+
}
90+
91+
private static SqlMediaType mediaTypeFromHeaders(RestRequest request) {
92+
String acceptType = getAcceptValue(request);
93+
SqlMediaType mediaType = acceptType != null
94+
? SqlMediaType.fromMediaTypeOrFormat(acceptType)
95+
: SqlMediaType.xContentType(request.getXContentType());
96+
return checkNonNullMediaType(mediaType, request);
97+
}
98+
99+
private static SqlMediaType mediaTypeFromParams(RestRequest request) {
100+
return SqlMediaType.fromMediaTypeOrFormat(request.param(URL_PARAM_FORMAT));
101+
}
102+
103+
104+
private static SqlMediaType validateColumnarRequest(boolean requestIsColumnar, SqlMediaType fromMediaType, RestRequest request) {
105+
if (requestIsColumnar && fromMediaType.isTextFormat()) {
106+
throw new IllegalArgumentException("Invalid use of [columnar] argument: cannot be used in combination with "
107+
+ "txt, csv or tsv formats");
108+
}
109+
return checkNonNullMediaType(fromMediaType, request);
110+
}
111+
112+
private static SqlMediaType checkNonNullMediaType(SqlMediaType mediaType, RestRequest request) {
113+
if (mediaType == null) {
114+
String msg = String.format(Locale.ROOT, "Invalid request content type: Accept=[%s], Content-Type=[%s], format=[%s]",
115+
request.header("Accept"), request.header("Content-Type"), request.param(URL_PARAM_FORMAT));
116+
throw new IllegalArgumentException(msg);
117+
}
118+
119+
return mediaType;
120+
}
121+
122+
// Partially lifted from https://github.com/elastic/elasticsearch/pull/64406 RestRequest#parseHeaderWithMediaType()
123+
private static @Nullable String getAcceptValue(RestRequest request) {
124+
Map<String, List<String>> headers = request.getHeaders();
125+
final String headerName = "Accept";
126+
127+
// TODO: make all usages of headers case-insensitive
128+
List<String> header = headers.get(headerName);
129+
if (header == null || header.isEmpty()) {
130+
return null;
131+
} else if (header.size() > 1) {
132+
throw new IllegalArgumentException("Incorrect header [" + headerName + "]. " +
133+
"Only one value should be provided");
134+
}
135+
String rawContentType = header.get(0);
136+
if (Strings.hasText(rawContentType)) {
137+
if ("*/*".equals(rawContentType)) {
138+
// */* means "I don't care" which we should treat like not specifying the header
139+
return null;
140+
}
141+
return rawContentType;
142+
} else {
143+
throw new IllegalArgumentException("Header [" + headerName + "] cannot be empty.");
144+
}
145+
146+
}
147+
}

‎x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/plugin/SqlMediaTypeParserTests.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.sql.plugin;
9+
10+
import org.elasticsearch.common.xcontent.NamedXContentRegistry;
11+
import org.elasticsearch.core.TimeValue;
12+
import org.elasticsearch.rest.RestRequest;
13+
import org.elasticsearch.test.ESTestCase;
14+
import org.elasticsearch.test.rest.FakeRestRequest;
15+
import org.elasticsearch.xpack.sql.action.SqlQueryRequest;
16+
import org.elasticsearch.xpack.sql.plugin.SqlMediaTypeParser.SqlMediaType;
17+
import org.elasticsearch.xpack.sql.proto.Mode;
18+
import org.elasticsearch.xpack.sql.proto.RequestInfo;
19+
20+
import java.util.Collections;
21+
import java.util.HashMap;
22+
import java.util.List;
23+
24+
import static org.elasticsearch.xpack.sql.plugin.SqlMediaTypeParser.getResponseMediaType;
25+
import static org.elasticsearch.xpack.sql.plugin.TextFormat.CSV;
26+
import static org.elasticsearch.xpack.sql.plugin.TextFormat.PLAIN_TEXT;
27+
import static org.elasticsearch.xpack.sql.plugin.TextFormat.TSV;
28+
import static org.elasticsearch.xpack.sql.proto.RequestInfo.CLIENT_IDS;
29+
import static org.hamcrest.CoreMatchers.is;
30+
31+
public class SqlMediaTypeParserTests extends ESTestCase {
32+
33+
public void testPlainTextDetection() {
34+
SqlMediaType text = getResponseMediaType(reqWithAccept("text/plain"), createTestInstance());
35+
assertThat(text.textFormat(), is(PLAIN_TEXT));
36+
}
37+
38+
public void testCsvDetection() {
39+
SqlMediaType text = getResponseMediaType(reqWithAccept("text/csv"), createTestInstance());
40+
assertThat(text.textFormat(), is(CSV));
41+
42+
text = getResponseMediaType(reqWithAccept("text/csv; delimiter=x"), createTestInstance());
43+
assertThat(text.textFormat(), is(CSV));
44+
}
45+
46+
public void testTsvDetection() {
47+
SqlMediaType text = getResponseMediaType(reqWithAccept("text/tab-separated-values"),
48+
createTestInstance());
49+
assertThat(text.textFormat(), is(TSV));
50+
}
51+
52+
public void testMediaTypeDetectionWithParameters() {
53+
assertThat(getResponseMediaType(reqWithAccept("text/plain; charset=utf-8"),
54+
createTestInstance()).textFormat(), is(PLAIN_TEXT));
55+
assertThat(getResponseMediaType(reqWithAccept("text/plain; header=present"),
56+
createTestInstance()).textFormat(), is(PLAIN_TEXT));
57+
assertThat(getResponseMediaType(reqWithAccept("text/plain; charset=utf-8; header=present"),
58+
createTestInstance()).textFormat(), is(PLAIN_TEXT));
59+
60+
assertThat(getResponseMediaType(reqWithAccept("text/csv; charset=utf-8"),
61+
createTestInstance()).textFormat(), is(CSV));
62+
assertThat(getResponseMediaType(reqWithAccept("text/csv; header=present"),
63+
createTestInstance()).textFormat(), is(CSV));
64+
assertThat(getResponseMediaType(reqWithAccept("text/csv; charset=utf-8; header=present"),
65+
createTestInstance()).textFormat(), is(CSV));
66+
67+
assertThat(getResponseMediaType(reqWithAccept("text/tab-separated-values; charset=utf-8"),
68+
createTestInstance()).textFormat(), is(TSV));
69+
assertThat(getResponseMediaType(reqWithAccept("text/tab-separated-values; header=present"),
70+
createTestInstance()).textFormat(), is(TSV));
71+
assertThat(getResponseMediaType(reqWithAccept("text/tab-separated-values; charset=utf-8; header=present"),
72+
createTestInstance()).textFormat(), is(TSV));
73+
}
74+
75+
public void testInvalidFormat() {
76+
IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
77+
() -> getResponseMediaType(reqWithAccept("text/garbage"), createTestInstance()));
78+
assertEquals(e.getMessage(), "invalid format [text/garbage]");
79+
}
80+
81+
public void testNoFormat() {
82+
IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
83+
() -> getResponseMediaType(new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY).build(),
84+
createTestInstance()));
85+
assertEquals(e.getMessage(), "Invalid request content type: Accept=[null], Content-Type=[null], format=[null]");
86+
}
87+
88+
private static RestRequest reqWithAccept(String acceptHeader) {
89+
90+
return new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY)
91+
.withHeaders(new HashMap<String, List<String>>() {{
92+
put("Content-Type", Collections.singletonList("application/json"));
93+
put("Accept", Collections.singletonList(acceptHeader));
94+
}}).build();
95+
}
96+
97+
protected SqlQueryRequest createTestInstance() {
98+
return new SqlQueryRequest(randomAlphaOfLength(10), Collections.emptyList(), null, null,
99+
randomZone(), between(1, Integer.MAX_VALUE), TimeValue.parseTimeValue(randomTimeValue(), null, "test"),
100+
TimeValue.parseTimeValue(randomTimeValue(), null, "test"), false, randomAlphaOfLength(10),
101+
new RequestInfo(Mode.PLAIN, randomFrom(randomFrom(CLIENT_IDS), randomAlphaOfLengthBetween(10, 20))),
102+
randomBoolean(), randomBoolean()).binaryCommunication(false);
103+
}
104+
}

0 commit comments

Comments
 (0)