Pass REST params and content to extensions

src/main/java/org/opensearch/sdk/ExtensionRestHandler.java

@@ -9,6 +9,7 @@
 
 import java.util.List;
 
+import org.opensearch.extensions.rest.ExtensionRestRequest;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.RestHandler;
 import org.opensearch.rest.RestHandler.Route;
@@ -29,12 +30,12 @@ public interface ExtensionRestHandler {
 
     /**
      * Handles REST Requests forwarded from OpenSearch for a configured route on an extension.
-     * Parameters are components of the {@link RestRequest} received from a user.
+     * Parameter contains components of the {@link RestRequest} received from a user.
      * This method corresponds to the {@link BaseRestHandler#prepareRequest} method.
      * As in that method, consumed parameters must be tracked and returned in the response.
      *
-     * @param restRequest a REST request object for a request to be forwarded to extensions
+     * @param request a REST request object for a request to be forwarded to extensions
      * @return An {@link ExtensionRestResponse} to the request.
      */
-    ExtensionRestResponse handleRequest(ExtensionRestRequest restRequest);
+    ExtensionRestResponse handleRequest(ExtensionRestRequest request);
 }

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

@@ -11,6 +11,7 @@
 
 import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.xcontent.XContentBuilder;
+import org.opensearch.extensions.rest.ExtensionRestRequest;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestStatus;
 
@@ -23,71 +24,76 @@ public class ExtensionRestResponse extends BytesRestResponse {
      * Key passed in {@link BytesRestResponse} headers to identify parameters consumed by the handler. For internal use.
      */
     static final String CONSUMED_PARAMS_KEY = "extension.consumed.parameters";
+    /**
+     * Key passed in {@link BytesRestResponse} headers to identify content consumed by the handler. For internal use.
+     */
+    static final String CONSUMED_CONTENT_KEY = "extension.consumed.content";
 
     /**
      * Creates a new response based on {@link XContentBuilder}.
      *
+     * @param request  the REST request being responded to.
      * @param status  The REST status.
      * @param builder  The builder for the response.
-     * @param consumedParams  Parameters consumed by the handler.
      */
-    public ExtensionRestResponse(RestStatus status, XContentBuilder builder, List<String> consumedParams) {
+    public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, XContentBuilder builder) {
         super(status, builder);
-        addConsumedParamHeader(consumedParams);
+        addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
     }
 
     /**
      * Creates a new plain text response.
      *
+     * @param request  the REST request being responded to.
      * @param status  The REST status.
      * @param content  A plain text response string.
-     * @param consumedParams  Parameters consumed by the handler.
      */
-    public ExtensionRestResponse(RestStatus status, String content, List<String> consumedParams) {
+    public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String content) {
         super(status, content);
-        addConsumedParamHeader(consumedParams);
+        addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
     }
 
     /**
      * Creates a new plain text response.
      *
+     * @param request  the REST request being responded to.
      * @param status  The REST status.
      * @param contentType  The content type of the response string.
      * @param content  A response string.
-     * @param consumedParams  Parameters consumed by the handler.
      */
-    public ExtensionRestResponse(RestStatus status, String contentType, String content, List<String> consumedParams) {
+    public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, String content) {
         super(status, contentType, content);
-        addConsumedParamHeader(consumedParams);
+        addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
     }
 
     /**
      * Creates a binary response.
      *
+     * @param request  the REST request being responded to.
      * @param status  The REST status.
      * @param contentType  The content type of the response bytes.
      * @param content  Response bytes.
-     * @param consumedParams  Parameters consumed by the handler.
      */
-    public ExtensionRestResponse(RestStatus status, String contentType, byte[] content, List<String> consumedParams) {
+    public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, byte[] content) {
         super(status, contentType, content);
-        addConsumedParamHeader(consumedParams);
+        addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
     }
 
     /**
      * Creates a binary response.
      *
+     * @param request  the REST request being responded to.
      * @param status  The REST status.
      * @param contentType  The content type of the response bytes.
      * @param content  Response bytes.
-     * @param consumedParams  Parameters consumed by the handler.
      */
-    public ExtensionRestResponse(RestStatus status, String contentType, BytesReference content, List<String> consumedParams) {
+    public ExtensionRestResponse(ExtensionRestRequest request, RestStatus status, String contentType, BytesReference content) {
         super(status, contentType, content);
-        addConsumedParamHeader(consumedParams);
+        addConsumedHeaders(request.consumedParams(), request.isContentConsumed());
     }
 
-    private void addConsumedParamHeader(List<String> consumedParams) {
+    private void addConsumedHeaders(List<String> consumedParams, boolean contentConusmed) {
         consumedParams.stream().forEach(p -> addHeader(CONSUMED_PARAMS_KEY, p));
+        addHeader(CONSUMED_CONTENT_KEY, Boolean.toString(contentConusmed));
     }
 }

src/main/java/org/opensearch/sdk/ExtensionsRunner.java

@@ -18,8 +18,8 @@
 import org.opensearch.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.common.io.stream.NamedWriteableRegistryParseRequest;
 import org.opensearch.extensions.OpenSearchRequest;
+import org.opensearch.extensions.rest.ExtensionRestRequest;
 import org.opensearch.extensions.rest.RegisterRestActionsRequest;
-import org.opensearch.extensions.rest.RestExecuteOnExtensionRequest;
 import org.opensearch.extensions.settings.RegisterCustomSettingsRequest;
 import org.opensearch.common.network.NetworkModule;
 import org.opensearch.common.network.NetworkService;
@@ -357,7 +357,7 @@ public void startTransportService(TransportService transportService) {
             ThreadPool.Names.GENERIC,
             false,
             false,
-            RestExecuteOnExtensionRequest::new,
+            ExtensionRestRequest::new,
             ((request, channel, task) -> channel.sendResponse(extensionsRestRequestHandler.handleRestExecuteOnExtensionRequest(request)))
         );
 
@@ -477,7 +477,7 @@ public void sendLocalNodeRequest(TransportService transportService) {
      * Requests the ActionListener onFailure method to be run by OpenSearch.  The result will be handled by a {@link ActionListenerOnFailureResponseHandler}.
      *
      * @param transportService  The TransportService defining the connection to OpenSearch.
-     * @param failureException The exception to be sent to OpenSearch
+     * @param failureException  The exception to be sent to OpenSearch
      */
     public void sendActionListenerOnFailureRequest(TransportService transportService, Exception failureException) {
         logger.info("Sending ActionListener onFailure request to OpenSearch");

src/main/java/org/opensearch/sdk/handlers/ExtensionsRestRequestHandler.java

@@ -10,13 +10,12 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.common.bytes.BytesReference;
-import org.opensearch.extensions.rest.RestExecuteOnExtensionRequest;
+import org.opensearch.extensions.rest.ExtensionRestRequest;
 import org.opensearch.extensions.rest.RestExecuteOnExtensionResponse;
 import org.opensearch.rest.RestStatus;
 import org.opensearch.sdk.ExtensionRestHandler;
 import org.opensearch.sdk.ExtensionsRunner;
 import org.opensearch.sdk.ExtensionRestPathRegistry;
-import org.opensearch.sdk.ExtensionRestRequest;
 import org.opensearch.sdk.ExtensionRestResponse;
 
 /**
@@ -33,24 +32,18 @@ public class ExtensionsRestRequestHandler {
      * @param request  The REST request to execute.
      * @return A response acknowledging the request.
      */
-    public RestExecuteOnExtensionResponse handleRestExecuteOnExtensionRequest(RestExecuteOnExtensionRequest request) {
+    public RestExecuteOnExtensionResponse handleRestExecuteOnExtensionRequest(ExtensionRestRequest request) {
 
-        ExtensionRestHandler restHandler = extensionRestPathRegistry.getHandler(request.getMethod(), request.getUri());
+        ExtensionRestHandler restHandler = extensionRestPathRegistry.getHandler(request.method(), request.uri());
         if (restHandler == null) {
             return new RestExecuteOnExtensionResponse(
                 RestStatus.NOT_FOUND,
-                "No handler for " + ExtensionRestPathRegistry.restPathToString(request.getMethod(), request.getUri())
+                "No handler for " + ExtensionRestPathRegistry.restPathToString(request.method(), request.uri())
             );
         }
-        // ExtensionRestRequest restRequest = new ExtensionRestRequest(request);
-        ExtensionRestRequest restRequest = new ExtensionRestRequest(
-            request.getMethod(),
-            request.getUri(),
-            request.getRequestIssuerIdentity()
-        );
 
         // Get response from extension
-        ExtensionRestResponse response = restHandler.handleRequest(restRequest);
+        ExtensionRestResponse response = restHandler.handleRequest(request);
         logger.info("Sending extension response to OpenSearch: " + response.status());
         return new RestExecuteOnExtensionResponse(
             response.status(),

src/main/java/org/opensearch/sdk/sample/helloworld/rest/RestHelloAction.java

@@ -7,15 +7,14 @@
  */
 package org.opensearch.sdk.sample.helloworld.rest;
 
+import org.opensearch.extensions.rest.ExtensionRestRequest;
 import org.opensearch.rest.RestHandler.Route;
 import org.opensearch.rest.RestRequest.Method;
 import org.opensearch.sdk.ExtensionRestHandler;
-import org.opensearch.sdk.ExtensionRestRequest;
 import org.opensearch.sdk.ExtensionRestResponse;
 
 import java.net.URLDecoder;
 import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
 import java.util.List;
 
 import static org.opensearch.rest.RestRequest.Method.GET;
@@ -39,31 +38,32 @@ public List<Route> routes() {
 
     @Override
     public ExtensionRestResponse handleRequest(ExtensionRestRequest request) {
-        // We need to track which parameters are consumed to pass back to OpenSearch
-        List<String> consumedParams = new ArrayList<>();
         Method method = request.method();
-        String uri = request.uri();
 
-        if (Method.GET.equals(method) && "/hello".equals(uri)) {
-            return new ExtensionRestResponse(OK, String.format(GREETING, worldName), consumedParams);
-        } else if (Method.PUT.equals(method) && uri.startsWith("/hello/")) {
-            // Placeholder code here for parameters in named wildcard paths
-            // Full implementation based on params() will be implemented as part of
-            // https://github.com/opensearch-project/opensearch-sdk-java/issues/111
-            String name = uri.substring("/hello/".length());
-            consumedParams.add("name");
-            try {
-                worldName = URLDecoder.decode(name, StandardCharsets.UTF_8);
-            } catch (IllegalArgumentException e) {
-                return new ExtensionRestResponse(BAD_REQUEST, e.getMessage(), consumedParams);
-            }
-            return new ExtensionRestResponse(OK, "Updated the world's name to " + worldName, consumedParams);
+        if (Method.GET.equals(method)) {
+            return handleGetRequest(request);
+        } else if (Method.PUT.equals(method)) {
+            return handlePutRequest(request);
         }
-        return new ExtensionRestResponse(
-            NOT_FOUND,
-            "Extension REST action improperly configured to handle " + method.name() + " " + uri,
-            consumedParams
-        );
+        return handleBadRequest(request);
+    }
+
+    private ExtensionRestResponse handleGetRequest(ExtensionRestRequest request) {
+        return new ExtensionRestResponse(request, OK, String.format(GREETING, worldName));
+    }
+
+    private ExtensionRestResponse handlePutRequest(ExtensionRestRequest request) {
+        String name = request.param("name");
+        try {
+            worldName = URLDecoder.decode(name, StandardCharsets.UTF_8);
+        } catch (IllegalArgumentException e) {
+            return new ExtensionRestResponse(request, BAD_REQUEST, e.getMessage());
+        }
+        return new ExtensionRestResponse(request, OK, "Updated the world's name to " + worldName);
+    }
+
+    private ExtensionRestResponse handleBadRequest(ExtensionRestRequest request) {
+        return new ExtensionRestResponse(request, NOT_FOUND, "Extension REST action improperly configured to handle " + request.toString());
     }
 
 }

src/test/java/org/opensearch/sdk/TestExtensionRestPathRegistry.java

@@ -4,6 +4,7 @@
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.opensearch.extensions.rest.ExtensionRestRequest;
 import org.opensearch.rest.RestHandler.Route;
 import org.opensearch.rest.RestRequest.Method;
 import org.opensearch.test.OpenSearchTestCase;