From c67a807906c3f6a67b2bff5bfa1d589a1a43915a Mon Sep 17 00:00:00 2001 From: Daniel Widdis Date: Thu, 12 Dec 2024 13:55:31 -0800 Subject: [PATCH] Pass sdkClient to IndicesHandler and EncryptorUtils classes Signed-off-by: Daniel Widdis --- .../flowframework/FlowFrameworkPlugin.java | 3 ++- .../indices/FlowFrameworkIndicesHandler.java | 5 +++++ .../flowframework/util/EncryptorUtils.java | 5 ++++- .../indices/FlowFrameworkIndicesHandlerTests.java | 15 ++++++++++++++- .../GetWorkflowTransportActionTests.java | 6 +++++- .../flowframework/util/EncryptorUtilsTests.java | 6 +++++- 6 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java index ac202ca3c..33a79b322 100644 --- a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java +++ b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java @@ -135,7 +135,6 @@ public Collection 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 createComponents( ) : Collections.emptyMap() ); + EncryptorUtils encryptorUtils = new EncryptorUtils(clusterService, client, sdkClient, xContentRegistry); FlowFrameworkIndicesHandler flowFrameworkIndicesHandler = new FlowFrameworkIndicesHandler( client, + sdkClient, clusterService, encryptorUtils, xContentRegistry diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java index 0ef0a7adb..0dd2d934e 100644 --- a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java +++ b/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 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()) { diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java index 02daeaff0..5748c7ae7 100644 --- a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java +++ b/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; } diff --git a/src/test/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandlerTests.java b/src/test/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandlerTests.java index a7dd7f75e..0701c6d54 100644 --- a/src/test/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandlerTests.java +++ b/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); diff --git a/src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java index 0ce92b3f0..07d53300b 100644 --- a/src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java +++ b/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, diff --git a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java index eb6ee1544..5bd9c333c 100644 --- a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java +++ b/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";