Skip to content

Commit

Permalink
Only send one extension info when initializing
Browse files Browse the repository at this point in the history 
Signed-off-by: Daniel Widdis <[email protected]>
  • Loading branch information
@dbwiddis
dbwiddis committed Aug 25, 2022
1 parent 7af4f07 commit 211567e
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 54 deletions.
52 changes: 19 additions & 33 deletions src/main/java/org/opensearch/sdk/ExtensionsRunner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,16 @@
*/
package org.opensearch.sdk;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
import static java.util.Collections.emptySet;
import static org.opensearch.common.UUIDs.randomBase64UUID;

import java.io.File;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
Expand All @@ -17,25 +25,22 @@
import org.opensearch.cluster.node.DiscoveryNode;
import org.opensearch.common.io.stream.NamedWriteableRegistry;
import org.opensearch.common.io.stream.NamedWriteableRegistryParseRequest;
import org.opensearch.extensions.OpenSearchRequest;
import org.opensearch.common.network.NetworkModule;
import org.opensearch.common.network.NetworkService;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.util.PageCacheRecycler;
import org.opensearch.extensions.DiscoveryExtension;
import org.opensearch.extensions.ExtensionBooleanResponse;
import org.opensearch.discovery.InitializeExtensionsRequest;
import org.opensearch.discovery.InitializeExtensionsResponse;
import org.opensearch.extensions.ExtensionBooleanResponse;
import org.opensearch.extensions.ExtensionRequest;
import org.opensearch.extensions.ExtensionsOrchestrator;
import org.opensearch.extensions.OpenSearchRequest;
import org.opensearch.extensions.RegisterRestActionsRequest;
import org.opensearch.index.IndicesModuleRequest;
import org.opensearch.index.IndicesModuleResponse;
import org.opensearch.indices.IndicesModule;
import org.opensearch.indices.breaker.CircuitBreakerService;
import org.opensearch.indices.breaker.NoneCircuitBreakerService;
import org.opensearch.transport.netty4.Netty4Transport;
import org.opensearch.transport.SharedGroupFactory;
import org.opensearch.sdk.handlers.ClusterSettingsResponseHandler;
import org.opensearch.sdk.handlers.ClusterStateResponseHandler;
import org.opensearch.sdk.handlers.LocalNodeResponseHandler;
Expand All @@ -44,21 +49,15 @@
import org.opensearch.threadpool.ThreadPool;
import org.opensearch.transport.ClusterConnectionManager;
import org.opensearch.transport.ConnectionManager;
import org.opensearch.transport.SharedGroupFactory;
import org.opensearch.transport.TransportInterceptor;
import org.opensearch.transport.TransportResponse;
import org.opensearch.transport.TransportService;
import org.opensearch.transport.TransportSettings;
import org.opensearch.transport.TransportResponse;

import java.io.File;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.opensearch.transport.netty4.Netty4Transport;

import static java.util.Collections.emptySet;
import static org.opensearch.common.UUIDs.randomBase64UUID;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;

/**
* The primary class to run an extension.
Expand Down Expand Up @@ -125,24 +124,11 @@ DiscoveryNode getOpensearchNode() {
*/
InitializeExtensionsResponse handleExtensionInitRequest(InitializeExtensionsRequest extensionInitRequest) {
logger.info("Registering Extension Request received from OpenSearch");
InitializeExtensionsResponse initializeExtensionsResponse = new InitializeExtensionsResponse(extensionSettings.getExtensionName());
opensearchNode = extensionInitRequest.getSourceNode();
// Fetch the unique ID
for (DiscoveryExtension de : extensionInitRequest.getExtensions()) {
if (de.getName().equals(extensionSettings.getExtensionName())) {
setUniqueId(de.getId());
break;
}
}
// We could avoid this check if we only send one extension in the initialize request, rather than the entire list
if (getUniqueId() == null) {
throw new IllegalArgumentException(
"Extension " + extensionSettings.getExtensionName() + " does not match any requested extension."
);
}
setUniqueId(extensionInitRequest.getExtension().getId());
// Successfully initialized. Send the response.
try {
return initializeExtensionsResponse;
return new InitializeExtensionsResponse(extensionSettings.getExtensionName());
} finally {
// After sending successful response to initialization, send the REST API
port = opensearchNode.getAddress().getPort();
Expand Down
39 changes: 18 additions & 21 deletions src/test/java/org/opensearch/sdk/TestExtensionsRunner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,22 +11,21 @@

package org.opensearch.sdk;

import static java.util.Collections.emptySet;
import static java.util.Collections.emptyMap;
import static java.util.Collections.emptySet;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
Expand All @@ -35,18 +34,18 @@
import org.opensearch.common.io.stream.NamedWriteableRegistryResponse;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.transport.TransportAddress;
import org.opensearch.extensions.DiscoveryExtension;
import org.opensearch.discovery.InitializeExtensionsRequest;
import org.opensearch.discovery.InitializeExtensionsResponse;
import org.opensearch.extensions.OpenSearchRequest;
import org.opensearch.extensions.DiscoveryExtension;
import org.opensearch.extensions.ExtensionsOrchestrator.OpenSearchRequestType;
import org.opensearch.extensions.OpenSearchRequest;
import org.opensearch.sdk.handlers.ClusterSettingsResponseHandler;
import org.opensearch.sdk.handlers.ClusterStateResponseHandler;
import org.opensearch.sdk.handlers.LocalNodeResponseHandler;
import org.opensearch.sdk.handlers.RegisterRestActionsResponseHandler;
import org.opensearch.test.OpenSearchTestCase;
import org.opensearch.transport.TransportService;
import org.opensearch.transport.Transport;
import org.opensearch.transport.TransportService;

public class TestExtensionsRunner extends OpenSearchTestCase {

Expand Down Expand Up @@ -107,25 +106,23 @@ public void testHandleExtensionInitRequest() throws UnknownHostException {
emptySet(),
Version.CURRENT
);
List<DiscoveryExtension> extensions = List.of(
new DiscoveryExtension(
EXTENSION_NAME,
"opensearch-sdk-1",
"",
"",
"",
sourceNode.getAddress(),
new HashMap<String, String>(),
null,
null
)
DiscoveryExtension extension = new DiscoveryExtension(
EXTENSION_NAME,
"opensearch-sdk-1",
"",
"",
"",
sourceNode.getAddress(),
new HashMap<String, String>(),
null,
null
);

// Set mocked transport service
extensionsRunner.setExtensionTransportService(this.transportService);
doNothing().when(this.transportService).connectToNode(sourceNode);

InitializeExtensionsRequest extensionInitRequest = new InitializeExtensionsRequest(sourceNode, extensions);
InitializeExtensionsRequest extensionInitRequest = new InitializeExtensionsRequest(sourceNode, extension);

InitializeExtensionsResponse response = extensionsRunner.handleExtensionInitRequest(extensionInitRequest);
// Test if name and unique ID are set
Expand Down

0 comments on commit 211567e

Please sign in to comment.