Pass sdkClient to IndicesHandler and EncryptorUtils classes

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

src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java

@@ -135,7 +135,6 @@ public Collection<Object> createComponents(
         Settings settings = environment.settings();
         flowFrameworkSettings = new FlowFrameworkSettings(clusterService, settings);
         MachineLearningNodeClient mlClient = new MachineLearningNodeClient(client);
-        EncryptorUtils encryptorUtils = new EncryptorUtils(clusterService, client, xContentRegistry);
         SdkClient sdkClient = SdkClientFactory.createSdkClient(
             client,
             xContentRegistry,
@@ -149,8 +148,10 @@ public Collection<Object> createComponents(
                 )
                 : Collections.emptyMap()
         );
+        EncryptorUtils encryptorUtils = new EncryptorUtils(clusterService, client, sdkClient, xContentRegistry);
         FlowFrameworkIndicesHandler flowFrameworkIndicesHandler = new FlowFrameworkIndicesHandler(
             client,
+            sdkClient,
             clusterService,
             encryptorUtils,
             xContentRegistry

src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java

@@ -49,6 +49,7 @@
 import org.opensearch.flowframework.util.ParseUtils;
 import org.opensearch.flowframework.workflow.WorkflowData;
 import org.opensearch.index.engine.VersionConflictEngineException;
+import org.opensearch.remote.metadata.client.SdkClient;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -79,6 +80,7 @@
 public class FlowFrameworkIndicesHandler {
     private static final Logger logger = LogManager.getLogger(FlowFrameworkIndicesHandler.class);
     private final Client client;
+    private final SdkClient sdkClient;
     private final ClusterService clusterService;
     private final EncryptorUtils encryptorUtils;
     private static final Map<String, AtomicBoolean> indexMappingUpdated = new HashMap<>();
@@ -90,17 +92,20 @@ public class FlowFrameworkIndicesHandler {
     /**
      * constructor
      * @param client the open search client
+     * @param sdkClient the remote metadata client
      * @param clusterService ClusterService
      * @param encryptorUtils encryption utility
      * @param xContentRegistry contentRegister to parse any response
      */
     public FlowFrameworkIndicesHandler(
         Client client,
+        SdkClient sdkClient,
         ClusterService clusterService,
         EncryptorUtils encryptorUtils,
         NamedXContentRegistry xContentRegistry
     ) {
         this.client = client;
+        this.sdkClient = sdkClient;
         this.clusterService = clusterService;
         this.encryptorUtils = encryptorUtils;
         for (FlowFrameworkIndex mlIndex : FlowFrameworkIndex.values()) {

src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java

@@ -30,6 +30,7 @@
 import org.opensearch.flowframework.model.Template;
 import org.opensearch.flowframework.model.Workflow;
 import org.opensearch.flowframework.model.WorkflowNode;
+import org.opensearch.remote.metadata.client.SdkClient;
 
 import javax.crypto.spec.SecretKeySpec;
 
@@ -68,6 +69,7 @@ public class EncryptorUtils {
 
     private final ClusterService clusterService;
     private final Client client;
+    private final SdkClient sdkClient;
     private String masterKey;
     private final NamedXContentRegistry xContentRegistry;
 
@@ -77,10 +79,11 @@ public class EncryptorUtils {
      * @param client the node client
      * @param xContentRegistry the OpenSearch XContent Registry
      */
-    public EncryptorUtils(ClusterService clusterService, Client client, NamedXContentRegistry xContentRegistry) {
+    public EncryptorUtils(ClusterService clusterService, Client client, SdkClient sdkClient, NamedXContentRegistry xContentRegistry) {
         this.masterKey = null;
         this.clusterService = clusterService;
         this.client = client;
+        this.sdkClient = sdkClient;
         this.xContentRegistry = xContentRegistry;
     }
 

src/test/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandlerTests.java

@@ -35,6 +35,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.flowframework.TestHelpers;
@@ -50,6 +51,8 @@
 import org.opensearch.flowframework.workflow.WorkflowData;
 import org.opensearch.index.engine.VersionConflictEngineException;
 import org.opensearch.index.get.GetResult;
+import org.opensearch.remote.metadata.client.SdkClient;
+import org.opensearch.remote.metadata.client.impl.SdkClientFactory;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
 
@@ -81,6 +84,7 @@
 public class FlowFrameworkIndicesHandlerTests extends OpenSearchTestCase {
     @Mock
     private Client client;
+    private SdkClient sdkClient;
     @Mock
     private CreateIndexStep createIndexStep;
     @Mock
@@ -94,6 +98,8 @@ public class FlowFrameworkIndicesHandlerTests extends OpenSearchTestCase {
     @Mock
     protected ClusterService clusterService;
     @Mock
+    protected NamedXContentRegistry namedXContentRegistry;
+    @Mock
     private FlowFrameworkIndicesHandler flowMock;
     private static final String META = "_meta";
     private static final String SCHEMA_VERSION_FIELD = "schemaVersion";
@@ -112,7 +118,14 @@ public void setUp() throws Exception {
         threadContext = new ThreadContext(settings);
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
-        flowFrameworkIndicesHandler = new FlowFrameworkIndicesHandler(client, clusterService, encryptorUtils, xContentRegistry());
+        sdkClient = SdkClientFactory.createSdkClient(client, namedXContentRegistry, Collections.emptyMap());
+        flowFrameworkIndicesHandler = new FlowFrameworkIndicesHandler(
+            client,
+            sdkClient,
+            clusterService,
+            encryptorUtils,
+            xContentRegistry()
+        );
         adminClient = mock(AdminClient.class);
         indicesAdminClient = mock(IndicesAdminClient.class);
         metadata = mock(Metadata.class);

src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java

@@ -31,6 +31,8 @@
 import org.opensearch.flowframework.model.WorkflowNode;
 import org.opensearch.flowframework.util.EncryptorUtils;
 import org.opensearch.index.get.GetResult;
+import org.opensearch.remote.metadata.client.SdkClient;
+import org.opensearch.remote.metadata.client.impl.SdkClientFactory;
 import org.opensearch.tasks.Task;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -56,6 +58,7 @@
 public class GetWorkflowTransportActionTests extends OpenSearchTestCase {
 
     private Client client;
+    private SdkClient sdkClient;
     private NamedXContentRegistry xContentRegistry;
     private GetWorkflowTransportAction getTemplateTransportAction;
     private FlowFrameworkIndicesHandler flowFrameworkIndicesHandler;
@@ -68,7 +71,8 @@ public void setUp() throws Exception {
         this.client = mock(Client.class);
         this.xContentRegistry = mock(NamedXContentRegistry.class);
         this.flowFrameworkIndicesHandler = mock(FlowFrameworkIndicesHandler.class);
-        this.encryptorUtils = new EncryptorUtils(mock(ClusterService.class), client, xContentRegistry);
+        this.sdkClient = SdkClientFactory.createSdkClient(client, xContentRegistry, Collections.emptyMap());
+        this.encryptorUtils = new EncryptorUtils(mock(ClusterService.class), client, sdkClient, xContentRegistry);
         ClusterService clusterService = mock(ClusterService.class);
         ClusterSettings clusterSettings = new ClusterSettings(
             Settings.EMPTY,

src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java

@@ -32,6 +32,8 @@
 import org.opensearch.flowframework.model.Template;
 import org.opensearch.flowframework.model.Workflow;
 import org.opensearch.flowframework.model.WorkflowNode;
+import org.opensearch.remote.metadata.client.SdkClient;
+import org.opensearch.remote.metadata.client.impl.SdkClientFactory;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
 
@@ -53,6 +55,7 @@ public class EncryptorUtilsTests extends OpenSearchTestCase {
 
     private ClusterService clusterService;
     private Client client;
+    private SdkClient sdkClient;
     private NamedXContentRegistry xContentRegistry;
     private EncryptorUtils encryptorUtils;
     private String testMasterKey;
@@ -66,7 +69,8 @@ public void setUp() throws Exception {
         this.clusterService = mock(ClusterService.class);
         this.client = mock(Client.class);
         this.xContentRegistry = mock(NamedXContentRegistry.class);
-        this.encryptorUtils = new EncryptorUtils(clusterService, client, xContentRegistry);
+        this.sdkClient = SdkClientFactory.createSdkClient(client, xContentRegistry, Collections.emptyMap());
+        this.encryptorUtils = new EncryptorUtils(clusterService, client, sdkClient, xContentRegistry);
         this.testMasterKey = encryptorUtils.generateMasterKey();
         this.testCredentialKey = "credential_key";
         this.testCredentialValue = "12345";