Merge branch 'main' into Issue-171

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Oct 10, 2022 · 1cb21b5 · 1cb21b5
2 parents b34714c + 34332b3
commit 1cb21b5
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 7 deletions.
diff --git a/src/main/java/org/opensearch/sdk/Extension.java b/src/main/java/org/opensearch/sdk/Extension.java
@@ -12,7 +12,11 @@
 import java.net.URL;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionResponse;
+import org.opensearch.action.support.TransportAction;
 import org.opensearch.common.settings.Setting;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -48,6 +52,15 @@ default List<Setting<?>> getSettings() {
         return Collections.emptyList();
     }
 
+    /**
+     * Gets an optional list of custom {@link TransportAction} for the extension to register with OpenSearch.
+     *
+     * @return a list of custom transport actions this extension uses.
+     */
+    default Map<String, Class<? extends TransportAction<ActionRequest, ActionResponse>>> getActions() {
+        return Collections.emptyMap();
+    }
+
     /**
      * Helper method to read extension settings from a YAML file.
      *

diff --git a/src/main/java/org/opensearch/sdk/ExtensionsRunner.java b/src/main/java/org/opensearch/sdk/ExtensionsRunner.java
@@ -49,7 +49,6 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
@@ -102,7 +101,7 @@ public class ExtensionsRunner {
     /**
      * Instantiates a new transportActions
      */
-    public TransportActions transportActions = new TransportActions(new HashMap<>());
+    public TransportActions transportActions;
 
     /**
      * Instantiates a new update settings request handler
@@ -130,6 +129,8 @@ protected ExtensionsRunner(Extension extension) throws IOException {
         }
         // save custom settings
         this.customSettings = extension.getSettings();
+        // save custom transport actions
+        this.transportActions = new TransportActions(extension.getActions());
         // initialize the transport service
         nettyTransport.initializeExtensionTransportService(this.getSettings(), this);
     }
@@ -150,7 +151,7 @@ public void setUniqueId(String id) {
         this.uniqueId = id;
     }
 
-    String getUniqueId() {
+    public String getUniqueId() {
         return uniqueId;
     }
 

diff --git a/src/main/java/org/opensearch/sdk/TransportActions.java b/src/main/java/org/opensearch/sdk/TransportActions.java
@@ -46,15 +46,16 @@ public <Request extends ActionRequest, Response extends ActionResponse> Transpor
      *
      * @param transportService  The TransportService defining the connection to OpenSearch.
      * @param opensearchNode The OpenSearch node where transport actions being registered.
+     * @param uniqueId The identity used to
      */
-    public void sendRegisterTransportActionsRequest(TransportService transportService, DiscoveryNode opensearchNode) {
+    public void sendRegisterTransportActionsRequest(TransportService transportService, DiscoveryNode opensearchNode, String uniqueId) {
         logger.info("Sending Register Transport Actions request to OpenSearch");
         ExtensionBooleanResponseHandler registerTransportActionsResponseHandler = new ExtensionBooleanResponseHandler();
         try {
             transportService.sendRequest(
                 opensearchNode,
                 ExtensionsOrchestrator.REQUEST_EXTENSION_REGISTER_TRANSPORT_ACTIONS,
-                new RegisterTransportActionsRequest("", transportActions),
+                new RegisterTransportActionsRequest(uniqueId, transportActions),
                 registerTransportActionsResponseHandler
             );
         } catch (Exception e) {

diff --git a/src/main/java/org/opensearch/sdk/handlers/ExtensionsInitRequestHandler.java b/src/main/java/org/opensearch/sdk/handlers/ExtensionsInitRequestHandler.java
@@ -48,7 +48,8 @@ public InitializeExtensionsResponse handleExtensionInitRequest(
             extensionsRunner.sendRegisterCustomSettingsRequest(extensionsRunner.extensionTransportService);
             extensionsRunner.transportActions.sendRegisterTransportActionsRequest(
                 extensionsRunner.extensionTransportService,
-                extensionsRunner.opensearchNode
+                extensionsRunner.opensearchNode,
+                extensionsRunner.getUniqueId()
             );
         }
     }

diff --git a/src/test/java/org/opensearch/sdk/TestExtensionTransportActionsAPI.java b/src/test/java/org/opensearch/sdk/TestExtensionTransportActionsAPI.java
@@ -82,7 +82,7 @@ public void setUp() throws Exception {
 
     @Test
     public void testRegisterTransportAction() {
-        transportActions.sendRegisterTransportActionsRequest(transportService, opensearchNode);
+        transportActions.sendRegisterTransportActionsRequest(transportService, opensearchNode, "test-extension");
         verify(transportService, times(1)).sendRequest(
             any(),
             eq(ExtensionsOrchestrator.REQUEST_EXTENSION_REGISTER_TRANSPORT_ACTIONS),