Pass full RestResponse to user from Extension (#114)

Signed-off-by: Daniel Widdis <[email protected]>

Signed-off-by: Daniel Widdis <[email protected]>

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);
 }

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()
+        );
     }
 
     /**

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<Route> 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
+        );
     }
 
 }

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

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"));
     }
 
 }