diff --git a/src/main/java/org/opensearch/sdk/ExtensionRestHandler.java b/src/main/java/org/opensearch/sdk/ExtensionRestHandler.java index fcb4ec33..9b2035a9 100644 --- a/src/main/java/org/opensearch/sdk/ExtensionRestHandler.java +++ b/src/main/java/org/opensearch/sdk/ExtensionRestHandler.java @@ -14,6 +14,7 @@ import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestRequest.Method; +import org.opensearch.rest.RestResponse; /** * This interface defines methods which an extension REST handler (action) must provide. @@ -33,7 +34,7 @@ public interface ExtensionRestHandler { * * @param method A REST method. * @param uri The URI to handle. - * @return A response string to be sent to the end user via OpenSearch. + * @return A {@link RestResponse} to the request. */ - String handleRequest(Method method, String uri); + RestResponse handleRequest(Method method, String uri); } diff --git a/src/main/java/org/opensearch/sdk/ExtensionsRunner.java b/src/main/java/org/opensearch/sdk/ExtensionsRunner.java index a0ab8363..aaa563ca 100644 --- a/src/main/java/org/opensearch/sdk/ExtensionsRunner.java +++ b/src/main/java/org/opensearch/sdk/ExtensionsRunner.java @@ -15,6 +15,7 @@ import org.opensearch.Version; import org.opensearch.cluster.ClusterModule; import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.common.bytes.BytesReference; import org.opensearch.common.io.stream.NamedWriteableRegistry; import org.opensearch.common.io.stream.NamedWriteableRegistryParseRequest; import org.opensearch.extensions.OpenSearchRequest; @@ -35,7 +36,9 @@ import org.opensearch.indices.IndicesModule; import org.opensearch.indices.breaker.CircuitBreakerService; import org.opensearch.indices.breaker.NoneCircuitBreakerService; +import org.opensearch.rest.RestStatus; import org.opensearch.rest.RestHandler.Route; +import org.opensearch.rest.RestResponse; import org.opensearch.transport.netty4.Netty4Transport; import org.opensearch.transport.SharedGroupFactory; import org.opensearch.sdk.handlers.ClusterSettingsResponseHandler; @@ -233,12 +236,17 @@ RestExecuteOnExtensionResponse handleRestExecuteOnExtensionRequest(RestExecuteOn String restPath = request.getMethod().name() + " " + request.getUri(); ExtensionRestHandler restHandler = extensionRestPathMap.get(restPath); if (restHandler == null) { - return new RestExecuteOnExtensionResponse("FAILED: No handler for " + restPath); + return new RestExecuteOnExtensionResponse(RestStatus.INTERNAL_SERVER_ERROR, "No handler for " + restPath); } // Get response from extension - String response = restHandler.handleRequest(request.getMethod(), request.getUri()); - logger.info("Sending extension response to OpenSearch: " + response); - return new RestExecuteOnExtensionResponse(response); + RestResponse response = restHandler.handleRequest(request.getMethod(), request.getUri()); + logger.info("Sending extension response to OpenSearch: " + response.status()); + return new RestExecuteOnExtensionResponse( + response.status(), + response.contentType(), + BytesReference.toBytes(response.content()), + response.getHeaders() + ); } /** diff --git a/src/main/java/org/opensearch/sdk/sample/rest/RestHelloAction.java b/src/main/java/org/opensearch/sdk/sample/rest/RestHelloAction.java index f37a7a37..667ca7e1 100644 --- a/src/main/java/org/opensearch/sdk/sample/rest/RestHelloAction.java +++ b/src/main/java/org/opensearch/sdk/sample/rest/RestHelloAction.java @@ -7,14 +7,18 @@ */ package org.opensearch.sdk.sample.rest; +import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest.Method; +import org.opensearch.rest.RestResponse; import org.opensearch.sdk.ExtensionRestHandler; import java.util.List; import static java.util.Collections.singletonList; import static org.opensearch.rest.RestRequest.Method.GET; +import static org.opensearch.rest.RestStatus.OK; +import static org.opensearch.rest.RestStatus.INTERNAL_SERVER_ERROR; /** * Sample REST Handler (REST Action). Extension REST handlers must implement {@link ExtensionRestHandler}. @@ -29,11 +33,14 @@ public List routes() { } @Override - public String handleRequest(Method method, String uri) { + public RestResponse handleRequest(Method method, String uri) { if (Method.GET.equals(method) && "/hello".equals(uri)) { - return GREETING; + return new BytesRestResponse(OK, GREETING); } - return null; + return new BytesRestResponse( + INTERNAL_SERVER_ERROR, + "Extension REST action improperly configured to handle " + method.name() + " " + uri + ); } } diff --git a/src/test/java/org/opensearch/sdk/TestExtensionsRunner.java b/src/test/java/org/opensearch/sdk/TestExtensionsRunner.java index 05410b62..30bd6c11 100644 --- a/src/test/java/org/opensearch/sdk/TestExtensionsRunner.java +++ b/src/test/java/org/opensearch/sdk/TestExtensionsRunner.java @@ -24,6 +24,7 @@ import java.net.InetAddress; import java.net.UnknownHostException; +import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.HashMap; @@ -41,7 +42,9 @@ import org.opensearch.extensions.OpenSearchRequest; import org.opensearch.extensions.rest.RestExecuteOnExtensionRequest; import org.opensearch.extensions.rest.RestExecuteOnExtensionResponse; +import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest.Method; +import org.opensearch.rest.RestStatus; import org.opensearch.sdk.handlers.ClusterSettingsResponseHandler; import org.opensearch.sdk.handlers.ClusterStateResponseHandler; import org.opensearch.sdk.handlers.LocalNodeResponseHandler; @@ -149,8 +152,12 @@ public void testHandleRestExecuteOnExtensionRequest() throws Exception { RestExecuteOnExtensionRequest request = new RestExecuteOnExtensionRequest(Method.GET, "/foo"); RestExecuteOnExtensionResponse response = extensionsRunner.handleRestExecuteOnExtensionRequest(request); - assertTrue(response.getResponse().contains("GET")); - assertTrue(response.getResponse().contains("/foo")); + // this will fail in test environment with no registered actions + assertEquals(RestStatus.INTERNAL_SERVER_ERROR, response.getStatus()); + assertEquals(BytesRestResponse.TEXT_CONTENT_TYPE, response.getContentType()); + String responseStr = new String(response.getContent(), StandardCharsets.UTF_8); + assertTrue(responseStr.contains("GET")); + assertTrue(responseStr.contains("/foo")); } @Test diff --git a/src/test/java/org/opensearch/sdk/sample/rest/TestRestHelloAction.java b/src/test/java/org/opensearch/sdk/sample/rest/TestRestHelloAction.java index 9462413b..5a888b32 100644 --- a/src/test/java/org/opensearch/sdk/sample/rest/TestRestHelloAction.java +++ b/src/test/java/org/opensearch/sdk/sample/rest/TestRestHelloAction.java @@ -7,12 +7,17 @@ */ package org.opensearch.sdk.sample.rest; +import java.nio.charset.StandardCharsets; import java.util.List; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest.Method; +import org.opensearch.common.bytes.BytesReference; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestResponse; +import org.opensearch.rest.RestStatus; import org.opensearch.sdk.ExtensionRestHandler; import org.opensearch.test.OpenSearchTestCase; @@ -37,9 +42,23 @@ public void testRoutes() { @Test public void testHandleRequest() { - assertEquals("Hello, World!", restHelloAction.handleRequest(Method.GET, "/hello")); - assertNull(restHelloAction.handleRequest(Method.PUT, "/hello")); - assertNull(restHelloAction.handleRequest(Method.GET, "/goodbye")); + RestResponse response = restHelloAction.handleRequest(Method.GET, "/hello"); + assertEquals(RestStatus.OK, response.status()); + assertEquals(BytesRestResponse.TEXT_CONTENT_TYPE, response.contentType()); + String responseStr = new String(BytesReference.toBytes(response.content()), StandardCharsets.UTF_8); + assertEquals("Hello, World!", responseStr); + + response = restHelloAction.handleRequest(Method.PUT, "/hello"); + assertEquals(RestStatus.INTERNAL_SERVER_ERROR, response.status()); + assertEquals(BytesRestResponse.TEXT_CONTENT_TYPE, response.contentType()); + responseStr = new String(BytesReference.toBytes(response.content()), StandardCharsets.UTF_8); + assertTrue(responseStr.contains("PUT")); + + response = restHelloAction.handleRequest(Method.GET, "/goodbye"); + assertEquals(RestStatus.INTERNAL_SERVER_ERROR, response.status()); + assertEquals(BytesRestResponse.TEXT_CONTENT_TYPE, response.contentType()); + responseStr = new String(BytesReference.toBytes(response.content()), StandardCharsets.UTF_8); + assertTrue(responseStr.contains("/goodbye")); } }