diff --git a/src/main/java/org/opensearch/sdk/Extension.java b/src/main/java/org/opensearch/sdk/Extension.java index 54ef2467..44fc053b 100644 --- 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> 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>> 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 index 3f0dae5b..b1114908 100644 --- a/src/main/java/org/opensearch/sdk/ExtensionsRunner.java +++ b/src/main/java/org/opensearch/sdk/ExtensionsRunner.java @@ -54,7 +54,6 @@ 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.Consumer; @@ -107,7 +106,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 @@ -127,6 +126,7 @@ public ExtensionsRunner() throws IOException { .put(TransportSettings.PORT.getKey(), extensionSettings.getHostPort()) .build(); this.customSettings = Collections.emptyList(); + this.transportActions = new TransportActions(Collections.emptyMap()); } /** @@ -150,6 +150,8 @@ private 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); // start listening on configured port and wait for connection from OpenSearch @@ -178,7 +180,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 index 49fc8ee9..e175a60c 100644 --- a/src/main/java/org/opensearch/sdk/TransportActions.java +++ b/src/main/java/org/opensearch/sdk/TransportActions.java @@ -46,15 +46,16 @@ public 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 index bd7f7c90..71503e49 100644 --- 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 index a998be53..b01d0c8a 100644 --- 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),