Register REST paths including named wildcards

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

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

@@ -0,0 +1,70 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.sdk;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.opensearch.common.path.PathTrie;
+import org.opensearch.rest.RestUtils;
+import org.opensearch.rest.RestRequest.Method;
+
+/**
+ * This class registers REST paths from extension Rest Handlers.
+ */
+public class ExtensionRestPathRegistry {
+
+    // PathTrie to match paths to handlers
+    private PathTrie<ExtensionRestHandler> pathTrie = new PathTrie<>(RestUtils.REST_DECODER);
+    // List to return registered handlers
+    private List<String> registeredPaths = new ArrayList<>();
+
+    /**
+     * Register a REST 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 extensionRestHandler  The RestHandler to handle this route
+     */
+    public void register(Method method, String uri, ExtensionRestHandler extensionRestHandler) {
+        String restPath = restPathToString(method, uri);
+        pathTrie.insert(restPath, extensionRestHandler);
+        registeredPaths.add(restPath);
+    }
+
+    /**
+     * Get the registered REST handler for the specified method and URI.
+     *
+     * @param method  the registered method.
+     * @param uri  the registered URI.
+     * @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));
+    }
+
+    /**
+     * List the registered routes.
+     *
+     * @return A list of strings identifying the registered routes.
+     */
+    public List<String> getRegisteredPaths() {
+        return registeredPaths;
+    }
+
+    /**
+     * Converts a REST method and URI to a string.
+     *
+     * @param method  the method.
+     * @param uri  the URI.
+     * @return A string appending the method and URI.
+     */
+    public static String restPathToString(Method method, String uri) {
+        return method.name() + " " + uri;
+    }
+}

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

@@ -56,11 +56,9 @@
 
 import java.io.File;
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -78,10 +76,10 @@ public class ExtensionsRunner {
     private static final Logger logger = LogManager.getLogger(ExtensionsRunner.class);
     private static final String NODE_NAME_SETTING = "node.name";
 
-    private Map<String, ExtensionRestHandler> extensionRestPathMap = new HashMap<>();
     private String uniqueId;
     private DiscoveryNode opensearchNode;
     private TransportService extensionTransportService = null;
+    private ExtensionRestPathRegistry extensionRestPathRegistry = new ExtensionRestPathRegistry();
 
     private final Settings settings;
     private final TransportInterceptor NOOP_TRANSPORT_INTERCEPTOR = new TransportInterceptor() {
@@ -123,8 +121,7 @@ private ExtensionsRunner(Extension extension) throws IOException {
         // store rest handlers in the map
         for (ExtensionRestHandler extensionRestHandler : extension.getExtensionRestHandlers()) {
             for (Route route : extensionRestHandler.routes()) {
-                String restPath = route.getMethod().name() + " " + route.getPath();
-                extensionRestPathMap.put(restPath, extensionRestHandler);
+                extensionRestPathRegistry.register(route.getMethod(), route.getPath(), extensionRestHandler);
             }
         }
         // initialize the transport service
@@ -233,10 +230,12 @@ ExtensionBooleanResponse handleIndicesModuleNameRequest(IndicesModuleRequest ind
      */
     RestExecuteOnExtensionResponse handleRestExecuteOnExtensionRequest(RestExecuteOnExtensionRequest request) {
 
-        String restPath = request.getMethod().name() + " " + request.getUri();
-        ExtensionRestHandler restHandler = extensionRestPathMap.get(restPath);
+        ExtensionRestHandler restHandler = extensionRestPathRegistry.getHandler(request.getMethod(), request.getUri());
         if (restHandler == null) {
-            return new RestExecuteOnExtensionResponse(RestStatus.INTERNAL_SERVER_ERROR, "No handler for " + restPath);
+            return new RestExecuteOnExtensionResponse(
+                RestStatus.NOT_FOUND,
+                "No handler for " + ExtensionRestPathRegistry.restPathToString(request.getMethod(), request.getUri())
+            );
         }
         // Get response from extension
         RestResponse response = restHandler.handleRequest(request.getMethod(), request.getUri());
@@ -401,7 +400,7 @@ public void startTransportService(TransportService transportService) {
      * @param transportService  The TransportService defining the connection to OpenSearch.
      */
     public void sendRegisterRestActionsRequest(TransportService transportService) {
-        List<String> extensionRestPaths = new ArrayList<>(extensionRestPathMap.keySet());
+        List<String> extensionRestPaths = extensionRestPathRegistry.getRegisteredPaths();
         logger.info("Sending Register REST Actions request to OpenSearch for " + extensionRestPaths);
         RegisterRestActionsResponseHandler registerActionsResponseHandler = new RegisterRestActionsResponseHandler();
         try {

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

@@ -0,0 +1,96 @@
+package org.opensearch.sdk;
+
+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.rest.RestResponse;
+import org.opensearch.test.OpenSearchTestCase;
+
+public class TestExtensionRestPathRegistry extends OpenSearchTestCase {
+
+    private ExtensionRestPathRegistry extensionRestPathRegistry = new ExtensionRestPathRegistry();
+
+    private ExtensionRestHandler fooHandler = new ExtensionRestHandler() {
+        @Override
+        public List<Route> routes() {
+            return List.of(new Route(Method.GET, "/foo"));
+        }
+
+        @Override
+        public RestResponse handleRequest(Method method, String uri) {
+            return null;
+        }
+    };
+    private ExtensionRestHandler barHandler = new ExtensionRestHandler() {
+        @Override
+        public List<Route> routes() {
+            return List.of(new Route(Method.PUT, "/bar/{planet}"));
+        }
+
+        @Override
+        public RestResponse handleRequest(Method method, String uri) {
+            return null;
+        }
+    };
+    private ExtensionRestHandler bazHandler = new ExtensionRestHandler() {
+        @Override
+        public List<Route> routes() {
+            return List.of(new Route(Method.POST, "/baz/{moon}/qux"), new Route(Method.PUT, "/bar/baz"));
+        }
+
+        @Override
+        public RestResponse handleRequest(Method method, String uri) {
+            return null;
+        }
+    };
+
+    @Override
+    @BeforeEach
+    public void setUp() throws Exception {
+        super.setUp();
+        for (ExtensionRestHandler handler : List.of(fooHandler, barHandler, bazHandler)) {
+            for (Route route : handler.routes()) {
+                extensionRestPathRegistry.register(route.getMethod(), route.getPath(), handler);
+            }
+        }
+    }
+
+    @Test
+    public void testRegisterConflicts() {
+        // Can't register same exact name
+        assertThrows(IllegalArgumentException.class, () -> extensionRestPathRegistry.register(Method.GET, "/foo", fooHandler));
+        // Can't register conflicting named wildcards
+        assertThrows(IllegalArgumentException.class, () -> extensionRestPathRegistry.register(Method.PUT, "/bar/{none}", barHandler));
+    }
+
+    @Test
+    public void testGetHandler() {
+        assertEquals(fooHandler, extensionRestPathRegistry.getHandler(Method.GET, "/foo"));
+        assertNull(extensionRestPathRegistry.getHandler(Method.PUT, "/foo"));
+
+        // Exact match and wildcard match can overlap, exact takes priority
+        assertEquals(barHandler, extensionRestPathRegistry.getHandler(Method.PUT, "/bar/mars"));
+        assertEquals(bazHandler, extensionRestPathRegistry.getHandler(Method.PUT, "/bar/baz"));
+        assertNull(extensionRestPathRegistry.getHandler(Method.PUT, "/bar/mars/bar"));
+
+        assertEquals(bazHandler, extensionRestPathRegistry.getHandler(Method.POST, "/baz/europa/qux"));
+        assertNull(extensionRestPathRegistry.getHandler(Method.POST, "/bar/europa"));
+    }
+
+    @Test
+    public void testGetRegisteredPaths() {
+        List<String> registeredPaths = extensionRestPathRegistry.getRegisteredPaths();
+        assertTrue(registeredPaths.contains("GET /foo"));
+        assertTrue(registeredPaths.contains("PUT /bar/{planet}"));
+        assertTrue(registeredPaths.contains("PUT /bar/baz"));
+        assertTrue(registeredPaths.contains("POST /baz/{moon}/qux"));
+    }
+
+    @Test
+    public void testRestPathToString() {
+        assertEquals("GET /foo", ExtensionRestPathRegistry.restPathToString(Method.GET, "/foo"));
+    }
+}

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

@@ -153,7 +153,7 @@ public void testHandleRestExecuteOnExtensionRequest() throws Exception {
         RestExecuteOnExtensionRequest request = new RestExecuteOnExtensionRequest(Method.GET, "/foo");
         RestExecuteOnExtensionResponse response = extensionsRunner.handleRestExecuteOnExtensionRequest(request);
         // this will fail in test environment with no registered actions
-        assertEquals(RestStatus.INTERNAL_SERVER_ERROR, response.getStatus());
+        assertEquals(RestStatus.NOT_FOUND, response.getStatus());
         assertEquals(BytesRestResponse.TEXT_CONTENT_TYPE, response.getContentType());
         String responseStr = new String(response.getContent(), StandardCharsets.UTF_8);
         assertTrue(responseStr.contains("GET"));