Skip to content

Commit 845d77d

Browse files
dbwiddisdbwiddis
committed
Move ExtensionRestResponse to OpenSearch
Signed-off-by: Daniel Widdis <widdis@gmail.com>
1 parent fafe32f commit 845d77d

3 files changed

Lines changed: 237 additions & 5 deletions

File tree

‎server/src/main/java/org/opensearch/extensions/rest/ExtensionRestResponse.java‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.extensions.rest;
10+
11+
import org.opensearch.common.bytes.BytesReference;
12+
import org.opensearch.common.xcontent.XContentBuilder;
13+
import org.opensearch.rest.BytesRestResponse;
14+
import org.opensearch.rest.RestStatus;
15+
import java.util.List;
16+
17+
/**
18+
* A subclass of {@link BytesRestResponse} which processes the consumed parameters into a custom header.
19+
*
20+
* @opensearch.api
21+
*/
22+
public class ExtensionRestResponse extends BytesRestResponse {
23+
24+
/**
25+
* Key passed in {@link BytesRestResponse} headers to identify parameters consumed by the handler. For internal use.
26+
*/
27+
static final String CONSUMED_PARAMS_KEY = "extension.consumed.parameters";
28+
/**
29+
* Key passed in {@link BytesRestResponse} headers to identify content consumed by the handler. For internal use.
30+
*/
31+
static final String CONSUMED_CONTENT_KEY = "extension.consumed.content";
32+
33+
/**
34+
* Creates a new response based on {@link XContentBuilder}.
35+
*
36+
* @param request the REST request being responded to.
37+
* @param status The REST status.
38+
* @param builder The builder for the response.
39+
*/
40+
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, XContentBuilder builder) {
41+
super(status, builder);
42+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
43+
}
44+
45+
/**
46+
* Creates a new plain text response.
47+
*
48+
* @param request the REST request being responded to.
49+
* @param status The REST status.
50+
* @param content A plain text response string.
51+
*/
52+
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String content) {
53+
super(status, content);
54+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
55+
}
56+
57+
/**
58+
* Creates a new plain text response.
59+
*
60+
* @param request the REST request being responded to.
61+
* @param status The REST status.
62+
* @param contentType The content type of the response string.
63+
* @param content A response string.
64+
*/
65+
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, String content) {
66+
super(status, contentType, content);
67+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
68+
}
69+
70+
/**
71+
* Creates a binary response.
72+
*
73+
* @param request the REST request being responded to.
74+
* @param status The REST status.
75+
* @param contentType The content type of the response bytes.
76+
* @param content Response bytes.
77+
*/
78+
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, byte[] content) {
79+
super(status, contentType, content);
80+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
81+
}
82+
83+
/**
84+
* Creates a binary response.
85+
*
86+
* @param request the REST request being responded to.
87+
* @param status The REST status.
88+
* @param contentType The content type of the response bytes.
89+
* @param content Response bytes.
90+
*/
91+
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, BytesReference content) {
92+
super(status, contentType, content);
93+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
94+
}
95+
96+
private void addConsumedHeaders(List<String> consumedParams, boolean contentConusmed) {
97+
consumedParams.stream().forEach(p -> addHeader(CONSUMED_PARAMS_KEY, p));
98+
addHeader(CONSUMED_CONTENT_KEY, Boolean.toString(contentConusmed));
99+
}
100+
}

‎server/src/main/java/org/opensearch/extensions/rest/RestSendToExtensionAction.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@
4141

4242
import static java.util.Collections.emptyMap;
4343
import static java.util.Collections.unmodifiableList;
44+
import static org.opensearch.extensions.rest.ExtensionRestResponse.CONSUMED_CONTENT_KEY;
45+
import static org.opensearch.extensions.rest.ExtensionRestResponse.CONSUMED_PARAMS_KEY;
4446

4547
/**
4648
* An action that forwards REST requests to an extension
@@ -49,8 +51,6 @@ public class RestSendToExtensionAction extends BaseRestHandler {
4951

5052
private static final String SEND_TO_EXTENSION_ACTION = "send_to_extension_action";
5153
private static final Logger logger = LogManager.getLogger(RestSendToExtensionAction.class);
52-
private static final String CONSUMED_PARAMS_KEY = "extension.consumed.parameters";
53-
private static final String CONSUMED_CONTENT_KEY = "extension.consumed.content";
5454
// To replace with user identity see https://github.com/opensearch-project/OpenSearch/pull/4247
5555
private static final Principal DEFAULT_PRINCIPAL = new Principal() {
5656
@Override
@@ -151,15 +151,16 @@ public void handleResponse(RestExecuteOnExtensionResponse response) {
151151
List<String> consumedContent = headers.get(CONSUMED_CONTENT_KEY);
152152
if (consumedContent != null) {
153153
// conditionally consume content
154-
if (consumedParams.stream().filter(c -> Boolean.parseBoolean(c)).count() > 0) {
154+
if (consumedContent.stream().filter(c -> Boolean.parseBoolean(c)).count() > 0) {
155155
request.content();
156156
}
157157
}
158-
Map<String, List<String>> headersWithoutConsumedParams = headers.entrySet()
158+
Map<String, List<String>> headersWithoutConsumedParamsOrContent = headers.entrySet()
159159
.stream()
160160
.filter(e -> !e.getKey().equals(CONSUMED_PARAMS_KEY))
161+
.filter(e -> !e.getKey().equals(CONSUMED_CONTENT_KEY))
161162
.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue()));
162-
restExecuteOnExtensionResponse.setHeaders(headersWithoutConsumedParams);
163+
restExecuteOnExtensionResponse.setHeaders(headersWithoutConsumedParamsOrContent);
163164
inProgressLatch.countDown();
164165
}
165166

‎server/src/test/java/org/opensearch/extensions/rest/ExtensionRestResponseTests.java‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
package org.opensearch.extensions.rest;
2+
3+
import java.io.IOException;
4+
import java.nio.ByteBuffer;
5+
import java.util.Collections;
6+
import java.util.List;
7+
8+
import org.opensearch.common.bytes.BytesArray;
9+
import org.opensearch.common.bytes.BytesReference;
10+
import org.opensearch.common.xcontent.XContentBuilder;
11+
import org.opensearch.common.xcontent.XContentType;
12+
import org.opensearch.rest.RestRequest.Method;
13+
import org.opensearch.test.OpenSearchTestCase;
14+
15+
import static org.opensearch.extensions.rest.ExtensionRestResponse.CONSUMED_PARAMS_KEY;
16+
import static org.opensearch.rest.BytesRestResponse.TEXT_CONTENT_TYPE;
17+
import static org.opensearch.rest.RestStatus.ACCEPTED;
18+
import static org.opensearch.rest.RestStatus.OK;
19+
20+
public class ExtensionRestResponseTests extends OpenSearchTestCase {
21+
22+
private static final String OCTET_CONTENT_TYPE = "application/octet-stream";
23+
private static final String JSON_CONTENT_TYPE = "application/json; charset=UTF-8";
24+
25+
private String testText;
26+
private byte[] testBytes;
27+
28+
@Override
29+
public void setUp() throws Exception {
30+
super.setUp();
31+
testText = "plain text";
32+
testBytes = new byte[] { 1, 2 };
33+
}
34+
35+
private ExtensionRestRequest generateTestRequest() {
36+
ExtensionRestRequest request = new ExtensionRestRequest(
37+
Method.GET,
38+
"/foo",
39+
Collections.emptyMap(),
40+
null,
41+
new BytesArray("Text Content"),
42+
null
43+
);
44+
// consume params "foo" and "bar"
45+
request.param("foo");
46+
request.param("bar");
47+
// consume content
48+
request.content();
49+
return request;
50+
}
51+
52+
public void testConstructorWithBuilder() throws IOException {
53+
XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent());
54+
builder.startObject();
55+
builder.field("status", ACCEPTED);
56+
builder.endObject();
57+
ExtensionRestRequest request = generateTestRequest();
58+
ExtensionRestResponse response = new ExtensionRestResponse(request, OK, builder);
59+
60+
assertEquals(OK, response.status());
61+
assertEquals(JSON_CONTENT_TYPE, response.contentType());
62+
assertEquals("{\"status\":\"ACCEPTED\"}", response.content().utf8ToString());
63+
List<String> consumedParams = response.getHeaders().get(CONSUMED_PARAMS_KEY);
64+
for (String param : consumedParams) {
65+
assertTrue(request.consumedParams().contains(param));
66+
}
67+
assertTrue(request.isContentConsumed());
68+
}
69+
70+
public void testConstructorWithPlainText() {
71+
ExtensionRestRequest request = generateTestRequest();
72+
ExtensionRestResponse response = new ExtensionRestResponse(request, OK, testText);
73+
74+
assertEquals(OK, response.status());
75+
assertEquals(TEXT_CONTENT_TYPE, response.contentType());
76+
assertEquals(testText, response.content().utf8ToString());
77+
List<String> consumedParams = response.getHeaders().get(CONSUMED_PARAMS_KEY);
78+
for (String param : consumedParams) {
79+
assertTrue(request.consumedParams().contains(param));
80+
}
81+
assertTrue(request.isContentConsumed());
82+
}
83+
84+
public void testConstructorWithText() {
85+
ExtensionRestRequest request = generateTestRequest();
86+
ExtensionRestResponse response = new ExtensionRestResponse(request, OK, TEXT_CONTENT_TYPE, testText);
87+
88+
assertEquals(OK, response.status());
89+
assertEquals(TEXT_CONTENT_TYPE, response.contentType());
90+
assertEquals(testText, response.content().utf8ToString());
91+
92+
List<String> consumedParams = response.getHeaders().get(CONSUMED_PARAMS_KEY);
93+
for (String param : consumedParams) {
94+
assertTrue(request.consumedParams().contains(param));
95+
}
96+
assertTrue(request.isContentConsumed());
97+
}
98+
99+
public void testConstructorWithByteArray() {
100+
ExtensionRestRequest request = generateTestRequest();
101+
ExtensionRestResponse response = new ExtensionRestResponse(request, OK, OCTET_CONTENT_TYPE, testBytes);
102+
103+
assertEquals(OK, response.status());
104+
assertEquals(OCTET_CONTENT_TYPE, response.contentType());
105+
assertArrayEquals(testBytes, BytesReference.toBytes(response.content()));
106+
List<String> consumedParams = response.getHeaders().get(CONSUMED_PARAMS_KEY);
107+
for (String param : consumedParams) {
108+
assertTrue(request.consumedParams().contains(param));
109+
}
110+
assertTrue(request.isContentConsumed());
111+
}
112+
113+
public void testConstructorWithBytesReference() {
114+
ExtensionRestRequest request = generateTestRequest();
115+
ExtensionRestResponse response = new ExtensionRestResponse(
116+
request,
117+
OK,
118+
OCTET_CONTENT_TYPE,
119+
BytesReference.fromByteBuffer(ByteBuffer.wrap(testBytes, 0, 2))
120+
);
121+
122+
assertEquals(OK, response.status());
123+
assertEquals(OCTET_CONTENT_TYPE, response.contentType());
124+
assertArrayEquals(testBytes, BytesReference.toBytes(response.content()));
125+
List<String> consumedParams = response.getHeaders().get(CONSUMED_PARAMS_KEY);
126+
for (String param : consumedParams) {
127+
assertTrue(request.consumedParams().contains(param));
128+
}
129+
assertTrue(request.isContentConsumed());
130+
}
131+
}

0 commit comments

Comments
 (0)