Pass REST params and content to extensions (opensearch-project#163)

* Rename/merge duplicate ExtensionRestRequest implementations Signed-off-by: Daniel Widdis <[email protected]> * Track consumed params inside request object Signed-off-by: Daniel Widdis <[email protected]> * Put back test for the "this will never happen" case Signed-off-by: Daniel Widdis <[email protected]> * Add content and expand param getters Signed-off-by: Daniel Widdis <[email protected]> * Rebase to main and update after conficts Signed-off-by: Daniel Widdis <[email protected]> * More merge conflict cleanup Signed-off-by: Daniel Widdis <[email protected]> * Whitespace diff reduction Signed-off-by: Daniel Widdis <[email protected]> * Change uri to path in RestRequest handling Signed-off-by: Daniel Widdis <[email protected]> Signed-off-by: Daniel Widdis <[email protected]>
kokibas · Mar 17, 2023 · f0a618e · f0a618e
1 parent bfaf46e
commit f0a618e
Show file tree

Hide file tree

Showing 12 changed files with 138 additions and 277 deletions.
diff --git a/src/main/java/org/opensearch/sdk/ExtensionRestHandler.java b/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,11 @@ 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);
 }
diff --git a/src/main/java/org/opensearch/sdk/ExtensionRestPathRegistry.java b/src/main/java/org/opensearch/sdk/ExtensionRestPathRegistry.java
@@ -28,11 +28,11 @@ public class ExtensionRestPathRegistry {
      * Register a REST handler to handle a method and route in this extension's path registry.
      *
      * @param method  The method to register.
-     * @param uri  The URI to register. May include named wildcards.
+     * @param path  The path to register. May include named wildcards.
      * @param extensionRestHandler  The RestHandler to handle this route
      */
-    public void registerHandler(Method method, String uri, ExtensionRestHandler extensionRestHandler) {
-        String restPath = restPathToString(method, uri);
+    public void registerHandler(Method method, String path, ExtensionRestHandler extensionRestHandler) {
+        String restPath = restPathToString(method, path);
         pathTrie.insert(restPath, extensionRestHandler);
         registeredPaths.add(restPath);
     }
@@ -41,11 +41,11 @@ public void registerHandler(Method method, String uri, ExtensionRestHandler exte
      * Get the registered REST handler for the specified method and URI.
      *
      * @param method  the registered method.
-     * @param uri  the registered URI.
+     * @param path  the registered path.
      * @return The REST handler registered to handle this method and URI combination if found, null otherwise.
      */
-    public ExtensionRestHandler getHandler(Method method, String uri) {
-        return pathTrie.retrieve(restPathToString(method, uri));
+    public ExtensionRestHandler getHandler(Method method, String path) {
+        return pathTrie.retrieve(restPathToString(method, path));
     }
 
     /**
@@ -58,13 +58,13 @@ public List<String> getRegisteredPaths() {
     }
 
     /**
-     * Converts a REST method and URI to a string.
+     * Converts a REST method and path to a space delimited string to be used as a map lookup key.
      *
      * @param method  the method.
-     * @param uri  the URI.
-     * @return A string appending the method and URI.
+     * @param path  the path.
+     * @return A string appending the method and path.
      */
-    public static String restPathToString(Method method, String uri) {
-        return method.name() + " " + uri;
+    public static String restPathToString(Method method, String path) {
+        return method.name() + " " + path;
     }
 }
diff --git a/src/main/java/org/opensearch/sdk/ExtensionRestRequest.java b/src/main/java/org/opensearch/sdk/ExtensionRestRequest.java
diff --git a/src/main/java/org/opensearch/sdk/ExtensionRestResponse.java b/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));
     }
 }
diff --git a/src/main/java/org/opensearch/sdk/ExtensionsRunner.java b/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");

diff --git a/src/main/java/org/opensearch/sdk/handlers/ExtensionsRestRequestHandler.java b/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.path());
         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.path())
             );
         }
-        // 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(),