Skip to content

Commit 159e6f2

Browse files
dbwiddisdbwiddis
committed
Add content and expand param getters
Signed-off-by: Daniel Widdis <widdis@gmail.com>
1 parent 2caf42f commit 159e6f2

4 files changed

Lines changed: 43 additions & 10 deletions

File tree

‎src/main/java/org/opensearch/sdk/ExtensionRestResponse.java‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,10 @@ public class ExtensionRestResponse extends BytesRestResponse {
2424
* Key passed in {@link BytesRestResponse} headers to identify parameters consumed by the handler. For internal use.
2525
*/
2626
static final String CONSUMED_PARAMS_KEY = "extension.consumed.parameters";
27+
/**
28+
* Key passed in {@link BytesRestResponse} headers to identify content consumed by the handler. For internal use.
29+
*/
30+
static final String CONSUMED_CONTENT_KEY = "extension.consumed.content";
2731

2832
/**
2933
* Creates a new response based on {@link XContentBuilder}.
@@ -34,7 +38,7 @@ public class ExtensionRestResponse extends BytesRestResponse {
3438
*/
3539
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, XContentBuilder builder) {
3640
super(status, builder);
37-
addConsumedParamHeader(request.consumedParams());
41+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
3842
}
3943

4044
/**
@@ -46,7 +50,7 @@ public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, XC
4650
*/
4751
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String content) {
4852
super(status, content);
49-
addConsumedParamHeader(request.consumedParams());
53+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
5054
}
5155

5256
/**
@@ -59,7 +63,7 @@ public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, St
5963
*/
6064
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, String content) {
6165
super(status, contentType, content);
62-
addConsumedParamHeader(request.consumedParams());
66+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
6367
}
6468

6569
/**
@@ -72,7 +76,7 @@ public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, St
7276
*/
7377
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, byte[] content) {
7478
super(status, contentType, content);
75-
addConsumedParamHeader(request.consumedParams());
79+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
7680
}
7781

7882
/**
@@ -85,10 +89,11 @@ public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, St
8589
*/
8690
public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, BytesReference content) {
8791
super(status, contentType, content);
88-
addConsumedParamHeader(request.consumedParams());
92+
addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
8993
}
9094

91-
private void addConsumedParamHeader(List<String> consumedParams) {
95+
private void addConsumedHeaders(List<String> consumedParams, boolean contentConusmed) {
9296
consumedParams.stream().forEach(p -> addHeader(CONSUMED_PARAMS_KEY, p));
97+
addHeader(CONSUMED_CONTENT_KEY, Boolean.toString(contentConusmed));
9398
}
9499
}

‎src/test/java/org/opensearch/sdk/TestExtensionRestResponse.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import org.junit.jupiter.api.BeforeEach;
99
import org.junit.jupiter.api.Test;
10+
import org.opensearch.common.bytes.BytesArray;
1011
import org.opensearch.common.bytes.BytesReference;
1112
import org.opensearch.common.xcontent.XContentBuilder;
1213
import org.opensearch.common.xcontent.XContentType;
@@ -34,10 +35,12 @@ public void setUp() throws Exception {
3435
super.setUp();
3536
testText = "plain text";
3637
testBytes = new byte[] { 1, 2 };
37-
request = new ExtensionRestRequest(Method.GET, "/foo", Collections.emptyMap(), null);
38+
request = new ExtensionRestRequest(Method.GET, "/foo", Collections.emptyMap(), null, new BytesArray("Text Content"), null);
3839
// consume params "foo" and "bar"
3940
request.param("foo");
4041
request.param("bar");
42+
// consume content
43+
request.content();
4144
}
4245

4346
@Test
@@ -55,6 +58,7 @@ public void testConstructorWithBuilder() throws IOException {
5558
for (String param : consumedParams) {
5659
assertTrue(request.consumedParams().contains(param));
5760
}
61+
assertTrue(request.isContentConsumed());
5862
}
5963

6064
@Test
@@ -68,6 +72,7 @@ public void testConstructorWithPlainText() {
6872
for (String param : consumedParams) {
6973
assertTrue(request.consumedParams().contains(param));
7074
}
75+
assertTrue(request.isContentConsumed());
7176
}
7277

7378
@Test
@@ -82,6 +87,7 @@ public void testConstructorWithText() {
8287
for (String param : consumedParams) {
8388
assertTrue(request.consumedParams().contains(param));
8489
}
90+
assertTrue(request.isContentConsumed());
8591
}
8692

8793
@Test
@@ -95,6 +101,7 @@ public void testConstructorWithByteArray() {
95101
for (String param : consumedParams) {
96102
assertTrue(request.consumedParams().contains(param));
97103
}
104+
assertTrue(request.isContentConsumed());
98105
}
99106

100107
@Test
@@ -113,5 +120,6 @@ public void testConstructorWithBytesReference() {
113120
for (String param : consumedParams) {
114121
assertTrue(request.consumedParams().contains(param));
115122
}
123+
assertTrue(request.isContentConsumed());
116124
}
117125
}

‎src/test/java/org/opensearch/sdk/TestExtensionsRunner.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import org.junit.jupiter.api.Test;
3636
import org.opensearch.Version;
3737
import org.opensearch.cluster.node.DiscoveryNode;
38+
import org.opensearch.common.bytes.BytesArray;
3839
import org.opensearch.common.io.stream.NamedWriteableRegistryResponse;
3940
import org.opensearch.common.settings.Settings;
4041
import org.opensearch.common.settings.Setting.Property;
@@ -167,6 +168,8 @@ public void testHandleExtensionRestRequest() throws Exception {
167168
Method.GET,
168169
"/foo",
169170
Collections.emptyMap(),
171+
null,
172+
new BytesArray("bar"),
170173
ext.generateToken(userPrincipal)
171174
);
172175
RestExecuteOnExtensionResponse response = extensionsRunner.handleRestExecuteOnExtensionRequest(request);

‎src/test/java/org/opensearch/sdk/sample/helloworld/rest/TestRestHelloAction.java‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import org.opensearch.identity.PrincipalIdentifierToken;
2020
import org.opensearch.rest.RestHandler.Route;
2121
import org.opensearch.rest.RestRequest.Method;
22+
import org.opensearch.common.bytes.BytesArray;
2223
import org.opensearch.common.bytes.BytesReference;
2324
import org.opensearch.extensions.rest.ExtensionRestRequest;
2425
import org.opensearch.rest.BytesRestResponse;
@@ -56,15 +57,31 @@ public void testHandleRequest() {
5657
PrincipalIdentifierToken token = extensionTokenProcessor.generateToken(userPrincipal);
5758
Map<String, String> params = Collections.emptyMap();
5859

59-
ExtensionRestRequest getRequest = new ExtensionRestRequest(Method.GET, "/hello", params, token);
60+
ExtensionRestRequest getRequest = new ExtensionRestRequest(Method.GET, "/hello", params, null, new BytesArray(""), token);
6061
ExtensionRestRequest putRequest = new ExtensionRestRequest(
6162
Method.PUT,
6263
"/hello/Passing+Test",
6364
Map.of("name", "Passing+Test"),
65+
null,
66+
new BytesArray(""),
67+
token
68+
);
69+
ExtensionRestRequest badRequest = new ExtensionRestRequest(
70+
Method.PUT,
71+
"/hello/Bad%Request",
72+
Map.of("name", "Bad%Request"),
73+
null,
74+
new BytesArray(""),
75+
token
76+
);
77+
ExtensionRestRequest unsuccessfulRequest = new ExtensionRestRequest(
78+
Method.POST,
79+
"/goodbye",
80+
params,
81+
null,
82+
new BytesArray(""),
6483
token
6584
);
66-
ExtensionRestRequest badRequest = new ExtensionRestRequest(Method.PUT, "/hello/Bad%Request", Map.of("name", "Bad%Request"), token);
67-
ExtensionRestRequest unsuccessfulRequest = new ExtensionRestRequest(Method.POST, "/goodbye", params, token);
6885

6986
RestResponse response = restHelloAction.handleRequest(getRequest);
7087
assertEquals(RestStatus.OK, response.status());

0 commit comments

Comments
 (0)