Pass around dynamicActionRegistry instead of ActionModule

Signed-off-by: Craig Perkins <cwperx@amazon.com>

server/src/main/java/org/opensearch/extensions/ExtensionsManager.java

@@ -230,7 +230,7 @@ private void registerRequestHandler(ActionModule actionModule) {
             false,
             RegisterRestActionsRequest::new,
             ((request, channel, task) -> channel.sendResponse(
-                restActionsRequestHandler.handleRegisterRestActionsRequest(request, actionModule)
+                restActionsRequestHandler.handleRegisterRestActionsRequest(request, actionModule.getDynamicActionRegistry())
             ))
         );
         transportService.registerRequestHandler(

server/src/main/java/org/opensearch/extensions/rest/RestActionsRequestHandler.java

@@ -8,7 +8,7 @@
 
 package org.opensearch.extensions.rest;
 
-import org.opensearch.action.ActionModule;
+import org.opensearch.action.ActionModule.DynamicActionRegistry;
 import org.opensearch.extensions.AcknowledgedResponse;
 import org.opensearch.extensions.DiscoveryExtensionNode;
 import org.opensearch.rest.RestController;
@@ -54,10 +54,17 @@ public RestActionsRequestHandler(
      * @return A {@link AcknowledgedResponse} indicating success.
      * @throws Exception if the request is not handled properly.
      */
-    public TransportResponse handleRegisterRestActionsRequest(RegisterRestActionsRequest restActionsRequest, ActionModule actionModule)
-        throws Exception {
+    public TransportResponse handleRegisterRestActionsRequest(
+        RegisterRestActionsRequest restActionsRequest,
+        DynamicActionRegistry dynamicActionRegistry
+    ) throws Exception {
         DiscoveryExtensionNode discoveryExtensionNode = extensionIdMap.get(restActionsRequest.getUniqueId());
-        RestHandler handler = new RestSendToExtensionAction(restActionsRequest, discoveryExtensionNode, transportService, actionModule);
+        RestHandler handler = new RestSendToExtensionAction(
+            restActionsRequest,
+            discoveryExtensionNode,
+            transportService,
+            dynamicActionRegistry
+        );
         restController.registerHandler(handler);
         return new AcknowledgedResponse(true);
     }

server/src/main/java/org/opensearch/rest/extensions/RestSendToExtensionAction.java

@@ -10,7 +10,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.opensearch.action.ActionModule;
+import org.opensearch.action.ActionModule.DynamicActionRegistry;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.io.stream.StreamInput;
@@ -85,7 +85,7 @@ public RestSendToExtensionAction(
         RegisterRestActionsRequest restActionsRequest,
         DiscoveryExtensionNode discoveryExtensionNode,
         TransportService transportService,
-        ActionModule actionModule
+        DynamicActionRegistry dynamicActionRegistry
     ) {
         this.pathPrefix = "/_extensions/_" + restActionsRequest.getUniqueId();
         RestRequest.Method method;
@@ -111,11 +111,11 @@ public RestSendToExtensionAction(
             if (name.isPresent()) {
                 NamedRoute nr = new NamedRoute(method, path, name.get());
                 restActionsAsRoutes.add(nr);
-                actionModule.getDynamicActionRegistry().registerDynamicRoute(nr, this);
+                dynamicActionRegistry.registerDynamicRoute(nr, this);
             } else {
                 Route r = new Route(method, path);
                 restActionsAsRoutes.add(r);
-                actionModule.getDynamicActionRegistry().registerDynamicRoute(r, this);
+                dynamicActionRegistry.registerDynamicRoute(r, this);
             }
         }
         this.routes = unmodifiableList(restActionsAsRoutes);

server/src/test/java/org/opensearch/extensions/ExtensionsManagerTests.java

@@ -97,6 +97,7 @@
 public class ExtensionsManagerTests extends OpenSearchTestCase {
     private TransportService transportService;
     private ActionModule actionModule;
+    private DynamicActionRegistry dynamicActionRegistry;
     private RestController restController;
     private SettingsModule settingsModule;
     private ClusterService clusterService;
@@ -161,6 +162,7 @@ public void setup() throws Exception {
             Collections.emptySet()
         );
         actionModule = mock(ActionModule.class);
+        dynamicActionRegistry = mock(DynamicActionRegistry.class);
         restController = new RestController(
             emptySet(),
             null,
@@ -476,7 +478,7 @@ public void testHandleRegisterRestActionsRequest() throws Exception {
         List<String> deprecatedActionsList = List.of("GET /deprecated/foo", "It's deprecated!");
         RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
         TransportResponse response = extensionsManager.getRestActionsRequestHandler()
-            .handleRegisterRestActionsRequest(registerActionsRequest, actionModule);
+            .handleRegisterRestActionsRequest(registerActionsRequest, actionModule.getDynamicActionRegistry());
         assertEquals(AcknowledgedResponse.class, response.getClass());
         assertTrue(((AcknowledgedResponse) response).getStatus());
     }
@@ -508,7 +510,8 @@ public void testHandleRegisterRestActionsRequestWithInvalidMethod() throws Excep
         RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
         expectThrows(
             IllegalArgumentException.class,
-            () -> extensionsManager.getRestActionsRequestHandler().handleRegisterRestActionsRequest(registerActionsRequest, actionModule)
+            () -> extensionsManager.getRestActionsRequestHandler()
+                .handleRegisterRestActionsRequest(registerActionsRequest, actionModule.getDynamicActionRegistry())
         );
     }
 
@@ -522,7 +525,8 @@ public void testHandleRegisterRestActionsRequestWithInvalidDeprecatedMethod() th
         RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
         expectThrows(
             IllegalArgumentException.class,
-            () -> extensionsManager.getRestActionsRequestHandler().handleRegisterRestActionsRequest(registerActionsRequest, actionModule)
+            () -> extensionsManager.getRestActionsRequestHandler()
+                .handleRegisterRestActionsRequest(registerActionsRequest, actionModule.getDynamicActionRegistry())
         );
     }
 
@@ -535,7 +539,8 @@ public void testHandleRegisterRestActionsRequestWithInvalidUri() throws Exceptio
         RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
         expectThrows(
             IllegalArgumentException.class,
-            () -> extensionsManager.getRestActionsRequestHandler().handleRegisterRestActionsRequest(registerActionsRequest, actionModule)
+            () -> extensionsManager.getRestActionsRequestHandler()
+                .handleRegisterRestActionsRequest(registerActionsRequest, dynamicActionRegistry)
         );
     }
 
@@ -548,7 +553,8 @@ public void testHandleRegisterRestActionsRequestWithInvalidDeprecatedUri() throw
         RegisterRestActionsRequest registerActionsRequest = new RegisterRestActionsRequest(uniqueIdStr, actionsList, deprecatedActionsList);
         expectThrows(
             IllegalArgumentException.class,
-            () -> extensionsManager.getRestActionsRequestHandler().handleRegisterRestActionsRequest(registerActionsRequest, actionModule)
+            () -> extensionsManager.getRestActionsRequestHandler()
+                .handleRegisterRestActionsRequest(registerActionsRequest, dynamicActionRegistry)
         );
     }
 

server/src/test/java/org/opensearch/extensions/rest/RestSendToExtensionActionTests.java

@@ -64,6 +64,7 @@ public class RestSendToExtensionActionTests extends OpenSearchTestCase {
     private MockNioTransport transport;
     private DiscoveryExtensionNode discoveryExtensionNode;
     private ActionModule actionModule;
+    private DynamicActionRegistry dynamicActionRegistry;
     private final ThreadPool threadPool = new TestThreadPool(RestSendToExtensionActionTests.class.getSimpleName());
 
     @Before
@@ -119,6 +120,7 @@ public void setup() throws Exception {
             null,
             new IdentityService(Settings.EMPTY, new ArrayList<>())
         );
+        dynamicActionRegistry = actionModule.getDynamicActionRegistry();
     }
 
     @Override
@@ -139,7 +141,7 @@ public void testRestSendToExtensionAction() throws Exception {
             registerRestActionRequest,
             discoveryExtensionNode,
             transportService,
-            actionModule
+            dynamicActionRegistry
         );
 
         assertEquals("send_to_extension_action", restSendToExtensionAction.getName());
@@ -171,7 +173,7 @@ public void testRestSendToExtensionActionWithNamedRoute() throws Exception {
             registerRestActionRequest,
             discoveryExtensionNode,
             transportService,
-            actionModule
+            dynamicActionRegistry
         );
 
         assertEquals("send_to_extension_action", restSendToExtensionAction.getName());
@@ -204,7 +206,7 @@ public void testRestSendToExtensionMultipleNamedRoutesWithSameName() throws Exce
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -216,7 +218,7 @@ public void testRestSendToExtensionMultipleRoutesWithSameMethodAndPath() throws
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -228,7 +230,7 @@ public void testRestSendToExtensionMultipleRoutesWithSameMethodAndPathWithDiffer
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -239,7 +241,7 @@ public void testRestSendToExtensionMultipleRoutesWithSameMethodAndPathWithPathPa
             List.of()
         );
         try {
-            new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule);
+            new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry);
         } catch (IllegalArgumentException e) {
             fail("IllegalArgumentException should not be thrown for different paths");
         }
@@ -261,7 +263,7 @@ public void testRestSendToExtensionWithNamedRouteCollidingWithDynamicTransportAc
 
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -275,7 +277,7 @@ public void testRestSendToExtensionWithNamedRouteCollidingWithNativeTransportAct
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -289,7 +291,7 @@ public void testRestSendToExtensionActionFilterHeaders() throws Exception {
             registerRestActionRequest,
             discoveryExtensionNode,
             transportService,
-            actionModule
+            dynamicActionRegistry
         );
 
         Map<String, List<String>> headers = new HashMap<>();
@@ -315,7 +317,7 @@ public void testRestSendToExtensionActionBadMethod() throws Exception {
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -327,7 +329,7 @@ public void testRestSendToExtensionActionBadDeprecatedMethod() throws Exception
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -339,7 +341,7 @@ public void testRestSendToExtensionActionMissingUri() throws Exception {
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 
@@ -351,7 +353,7 @@ public void testRestSendToExtensionActionMissingDeprecatedUri() throws Exception
         );
         expectThrows(
             IllegalArgumentException.class,
-            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, actionModule)
+            () -> new RestSendToExtensionAction(registerRestActionRequest, discoveryExtensionNode, transportService, dynamicActionRegistry)
         );
     }
 }