Skip to content

Commit 1375ee6

Browse files
siddharthabingiSiddhartha Bingi
siddharthabingi
and
Siddhartha Bingi
authored
Passing tenantId to MLClient for Connector, Model and ModelGroup (#1020)
Signed-off-by: Siddhartha Bingi <[email protected]> Co-authored-by: Siddhartha Bingi <[email protected]>
1 parent 1c16b1b commit 1375ee6

File tree

7 files changed

+23
-18
lines changed

7 files changed

+23
-18
lines changed

‎src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java

+1
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ public void onFailure(Exception ex) {
154154
.parameters(parameters)
155155
.credential(credentials)
156156
.actions(actions)
157+
.tenantId(tenantId)
157158
.build();
158159

159160
mlClient.createConnector(mlInput, actionListener);

‎src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ public PlainActionFuture<WorkflowData> execute(
7272
);
7373
String connectorId = (String) inputs.get(CONNECTOR_ID);
7474

75-
mlClient.deleteConnector(connectorId, new ActionListener<>() {
75+
mlClient.deleteConnector(connectorId, tenantId, new ActionListener<>() {
7676
@Override
7777
public void onResponse(DeleteResponse deleteResponse) {
7878
deleteConnectorFuture.onResponse(

‎src/main/java/org/opensearch/flowframework/workflow/DeleteModelStep.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ public PlainActionFuture<WorkflowData> execute(
7575

7676
String modelId = inputs.get(MODEL_ID).toString();
7777

78-
mlClient.deleteModel(modelId, new ActionListener<>() {
78+
mlClient.deleteModel(modelId, tenantId, new ActionListener<>() {
7979
@Override
8080
public void onResponse(DeleteResponse deleteResponse) {
8181
deleteModelFuture.onResponse(

‎src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java

+1
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ public void onFailure(Exception ex) {
125125

126126
MLRegisterModelGroupInputBuilder builder = MLRegisterModelGroupInput.builder();
127127
builder.name(modelGroupName);
128+
builder.tenantId(tenantId);
128129
if (description != null) {
129130
builder.description(description);
130131
}

‎src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,8 @@ public PlainActionFuture<WorkflowData> execute(
104104
MLRegisterModelInputBuilder builder = MLRegisterModelInput.builder()
105105
.functionName(FunctionName.REMOTE)
106106
.modelName(modelName)
107-
.connectorId(connectorId);
107+
.connectorId(connectorId)
108+
.tenantId(tenantId);
108109

109110
if (modelGroupId != null) {
110111
builder.modelGroupId(modelGroupId);

‎src/test/java/org/opensearch/flowframework/workflow/DeleteConnectorStepTests.java

+7-6
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import static org.opensearch.flowframework.common.WorkflowResources.CONNECTOR_ID;
3030
import static org.mockito.ArgumentMatchers.any;
3131
import static org.mockito.ArgumentMatchers.anyString;
32+
import static org.mockito.ArgumentMatchers.nullable;
3233
import static org.mockito.Mockito.doAnswer;
3334
import static org.mockito.Mockito.verify;
3435

@@ -54,12 +55,12 @@ public void testDeleteConnector() throws IOException, ExecutionException, Interr
5455

5556
doAnswer(invocation -> {
5657
String connectorIdArg = invocation.getArgument(0);
57-
ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
58+
ActionListener<DeleteResponse> actionListener = invocation.getArgument(2);
5859
ShardId shardId = new ShardId(new Index("indexName", "uuid"), 1);
5960
DeleteResponse output = new DeleteResponse(shardId, connectorIdArg, 1, 1, 1, true);
6061
actionListener.onResponse(output);
6162
return null;
62-
}).when(machineLearningNodeClient).deleteConnector(anyString(), anyActionListener());
63+
}).when(machineLearningNodeClient).deleteConnector(anyString(), nullable(String.class), anyActionListener());
6364

6465
PlainActionFuture<WorkflowData> future = deleteConnectorStep.execute(
6566
inputData.getNodeId(),
@@ -69,7 +70,7 @@ public void testDeleteConnector() throws IOException, ExecutionException, Interr
6970
Collections.emptyMap(),
7071
null
7172
);
72-
verify(machineLearningNodeClient).deleteConnector(anyString(), anyActionListener());
73+
verify(machineLearningNodeClient).deleteConnector(anyString(), nullable(String.class), anyActionListener());
7374

7475
assertTrue(future.isDone());
7576
assertEquals(connectorId, future.get().getContent().get(CONNECTOR_ID));
@@ -97,10 +98,10 @@ public void testDeleteConnectorFailure() throws IOException {
9798
DeleteConnectorStep deleteConnectorStep = new DeleteConnectorStep(machineLearningNodeClient);
9899

99100
doAnswer(invocation -> {
100-
ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
101+
ActionListener<DeleteResponse> actionListener = invocation.getArgument(2);
101102
actionListener.onFailure(new FlowFrameworkException("Failed to delete connector", RestStatus.INTERNAL_SERVER_ERROR));
102103
return null;
103-
}).when(machineLearningNodeClient).deleteConnector(anyString(), anyActionListener());
104+
}).when(machineLearningNodeClient).deleteConnector(anyString(), nullable(String.class), anyActionListener());
104105

105106
PlainActionFuture<WorkflowData> future = deleteConnectorStep.execute(
106107
inputData.getNodeId(),
@@ -111,7 +112,7 @@ public void testDeleteConnectorFailure() throws IOException {
111112
null
112113
);
113114

114-
verify(machineLearningNodeClient).deleteConnector(anyString(), anyActionListener());
115+
verify(machineLearningNodeClient).deleteConnector(anyString(), nullable(String.class), anyActionListener());
115116

116117
assertTrue(future.isDone());
117118
ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());

‎src/test/java/org/opensearch/flowframework/workflow/DeleteModelStepTests.java

+10-9
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929

3030
import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
3131
import static org.mockito.ArgumentMatchers.any;
32+
import static org.mockito.ArgumentMatchers.nullable;
3233
import static org.mockito.Mockito.doAnswer;
3334
import static org.mockito.Mockito.verify;
3435

@@ -54,12 +55,12 @@ public void testDeleteModel() throws IOException, ExecutionException, Interrupte
5455

5556
doAnswer(invocation -> {
5657
String modelIdArg = invocation.getArgument(0);
57-
ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
58+
ActionListener<DeleteResponse> actionListener = invocation.getArgument(2);
5859
ShardId shardId = new ShardId(new Index("indexName", "uuid"), 1);
5960
DeleteResponse output = new DeleteResponse(shardId, modelIdArg, 1, 1, 1, true);
6061
actionListener.onResponse(output);
6162
return null;
62-
}).when(machineLearningNodeClient).deleteModel(any(String.class), any());
63+
}).when(machineLearningNodeClient).deleteModel(any(String.class), nullable(String.class), any());
6364

6465
PlainActionFuture<WorkflowData> future = deleteModelStep.execute(
6566
inputData.getNodeId(),
@@ -69,7 +70,7 @@ public void testDeleteModel() throws IOException, ExecutionException, Interrupte
6970
Collections.emptyMap(),
7071
null
7172
);
72-
verify(machineLearningNodeClient).deleteModel(any(String.class), any());
73+
verify(machineLearningNodeClient).deleteModel(any(String.class), nullable(String.class), any());
7374

7475
assertTrue(future.isDone());
7576
assertEquals(modelId, future.get().getContent().get(MODEL_ID));
@@ -81,10 +82,10 @@ public void testDeleteModelNotFound() throws IOException, ExecutionException, In
8182
DeleteModelStep deleteModelStep = new DeleteModelStep(machineLearningNodeClient);
8283

8384
doAnswer(invocation -> {
84-
ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
85+
ActionListener<DeleteResponse> actionListener = invocation.getArgument(2);
8586
actionListener.onFailure(new OpenSearchStatusException("No model found with that id", RestStatus.NOT_FOUND));
8687
return null;
87-
}).when(machineLearningNodeClient).deleteModel(any(String.class), any());
88+
}).when(machineLearningNodeClient).deleteModel(any(String.class), nullable(String.class), any());
8889

8990
PlainActionFuture<WorkflowData> future = deleteModelStep.execute(
9091
inputData.getNodeId(),
@@ -94,7 +95,7 @@ public void testDeleteModelNotFound() throws IOException, ExecutionException, In
9495
Collections.emptyMap(),
9596
null
9697
);
97-
verify(machineLearningNodeClient).deleteModel(any(String.class), any());
98+
verify(machineLearningNodeClient).deleteModel(any(String.class), nullable(String.class), any());
9899

99100
assertTrue(future.isDone());
100101
assertEquals(modelId, future.get().getContent().get(MODEL_ID));
@@ -122,10 +123,10 @@ public void testDeleteModelFailure() throws IOException {
122123
DeleteModelStep deleteModelStep = new DeleteModelStep(machineLearningNodeClient);
123124

124125
doAnswer(invocation -> {
125-
ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
126+
ActionListener<DeleteResponse> actionListener = invocation.getArgument(2);
126127
actionListener.onFailure(new FlowFrameworkException("Failed to delete model", RestStatus.INTERNAL_SERVER_ERROR));
127128
return null;
128-
}).when(machineLearningNodeClient).deleteModel(any(String.class), any());
129+
}).when(machineLearningNodeClient).deleteModel(any(String.class), nullable(String.class), any());
129130

130131
PlainActionFuture<WorkflowData> future = deleteModelStep.execute(
131132
inputData.getNodeId(),
@@ -136,7 +137,7 @@ public void testDeleteModelFailure() throws IOException {
136137
null
137138
);
138139

139-
verify(machineLearningNodeClient).deleteModel(any(String.class), any());
140+
verify(machineLearningNodeClient).deleteModel(any(String.class), nullable(String.class), any());
140141

141142
assertTrue(future.isDone());
142143
ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());

0 commit comments

Comments
 (0)