Skip to content

Commit 61012a5

Browse files
opensearch-trigger-bot[bot]ylwu-amzn
opensearch-trigger-bot[bot]
and
ylwu-amzn
authored
fix update connector API (#1484) (#1494)
* fix update connector API Signed-off-by: Yaliang Wu <[email protected]> * fix ut Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> (cherry picked from commit 2f3f39e) Co-authored-by: Yaliang Wu <[email protected]>
1 parent cea1cd6 commit 61012a5

File tree

13 files changed

+275
-133
lines changed

13 files changed

+275
-133
lines changed

‎common/src/main/java/org/opensearch/ml/common/connector/Connector.java‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -30,17 +30,7 @@
3030
import org.opensearch.ml.common.AccessMode;
3131
import org.opensearch.ml.common.MLCommonsClassLoader;
3232
import org.opensearch.ml.common.output.model.ModelTensor;
33-
34-
import java.io.IOException;
35-
import java.security.AccessController;
36-
import java.security.PrivilegedActionException;
37-
import java.security.PrivilegedExceptionAction;
38-
import java.util.List;
39-
import java.util.Map;
40-
import java.util.Optional;
41-
import java.util.function.Function;
42-
import java.util.regex.Matcher;
43-
import java.util.regex.Pattern;
33+
import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
4434

4535
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
4636
import static org.opensearch.ml.common.utils.StringUtils.gson;
@@ -80,6 +70,7 @@ public interface Connector extends ToXContentObject, Writeable {
8070

8171
void writeTo(StreamOutput out) throws IOException;
8272

73+
void update(MLCreateConnectorInput updateContent, Function<String, String> function);
8374

8475
<T> void parseResponse(T orElse, List<ModelTensor> modelTensors, boolean b) throws IOException;
8576

‎common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.opensearch.core.xcontent.XContentBuilder;
3535
import org.opensearch.core.xcontent.XContentParser;
3636
import org.opensearch.ml.common.AccessMode;
37+
import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
3738

3839
@Log4j2
3940
@NoArgsConstructor
@@ -248,6 +249,38 @@ public void writeTo(StreamOutput out) throws IOException {
248249
}
249250
}
250251

252+
@Override
253+
public void update(MLCreateConnectorInput updateContent, Function<String, String> function) {
254+
if (updateContent.getName() != null) {
255+
this.name = updateContent.getName();
256+
}
257+
if (updateContent.getDescription() != null) {
258+
this.description = updateContent.getDescription();
259+
}
260+
if (updateContent.getVersion() != null) {
261+
this.version = updateContent.getVersion();
262+
}
263+
if (updateContent.getProtocol() != null) {
264+
this.protocol = updateContent.getProtocol();
265+
}
266+
if (updateContent.getParameters() != null && updateContent.getParameters().size() > 0) {
267+
this.parameters = updateContent.getParameters();
268+
}
269+
if (updateContent.getCredential() != null && updateContent.getCredential().size() > 0) {
270+
this.credential = updateContent.getCredential();
271+
encrypt(function);
272+
}
273+
if (updateContent.getActions() != null) {
274+
this.actions = updateContent.getActions();
275+
}
276+
if (updateContent.getBackendRoles() != null) {
277+
this.backendRoles = updateContent.getBackendRoles();
278+
}
279+
if (updateContent.getAccess() != null) {
280+
this.access = updateContent.getAccess();
281+
}
282+
}
283+
251284
@Override
252285
public <T> T createPredictPayload(Map<String, String> parameters) {
253286
Optional<ConnectorAction> predictAction = findPredictAction();

‎common/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ public class MLCreateConnectorInput implements ToXContentObject, Writeable {
5656
private Boolean addAllBackendRoles;
5757
private AccessMode access;
5858
private boolean dryRun = false;
59+
private boolean updateConnector = false;
5960

6061
@Builder(toBuilder = true)
6162
public MLCreateConnectorInput(String name,
@@ -68,9 +69,10 @@ public MLCreateConnectorInput(String name,
6869
List<String> backendRoles,
6970
Boolean addAllBackendRoles,
7071
AccessMode access,
71-
boolean dryRun
72+
boolean dryRun,
73+
boolean updateConnector
7274
) {
73-
if (!dryRun) {
75+
if (!dryRun && !updateConnector) {
7476
if (name == null) {
7577
throw new IllegalArgumentException("Connector name is null");
7678
}
@@ -92,9 +94,14 @@ public MLCreateConnectorInput(String name,
9294
this.addAllBackendRoles = addAllBackendRoles;
9395
this.access = access;
9496
this.dryRun = dryRun;
97+
this.updateConnector = updateConnector;
9598
}
9699

97100
public static MLCreateConnectorInput parse(XContentParser parser) throws IOException {
101+
return parse(parser, false);
102+
}
103+
104+
public static MLCreateConnectorInput parse(XContentParser parser, boolean updateConnector) throws IOException {
98105
String name = null;
99106
String description = null;
100107
String version = null;
@@ -159,7 +166,7 @@ public static MLCreateConnectorInput parse(XContentParser parser) throws IOExcep
159166
break;
160167
}
161168
}
162-
return new MLCreateConnectorInput(name, description, version, protocol, parameters, credential, actions, backendRoles, addAllBackendRoles, access, dryRun);
169+
return new MLCreateConnectorInput(name, description, version, protocol, parameters, credential, actions, backendRoles, addAllBackendRoles, access, dryRun, updateConnector);
163170
}
164171

165172
@Override
@@ -201,10 +208,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
201208

202209
@Override
203210
public void writeTo(StreamOutput output) throws IOException {
204-
output.writeString(name);
211+
output.writeOptionalString(name);
205212
output.writeOptionalString(description);
206-
output.writeString(version);
207-
output.writeString(protocol);
213+
output.writeOptionalString(version);
214+
output.writeOptionalString(protocol);
208215
if (parameters != null) {
209216
output.writeBoolean(true);
210217
output.writeMap(parameters, StreamOutput::writeString, StreamOutput::writeString);
@@ -240,13 +247,14 @@ public void writeTo(StreamOutput output) throws IOException {
240247
output.writeBoolean(false);
241248
}
242249
output.writeBoolean(dryRun);
250+
output.writeBoolean(updateConnector);
243251
}
244252

245253
public MLCreateConnectorInput(StreamInput input) throws IOException {
246-
name = input.readString();
254+
name = input.readOptionalString();
247255
description = input.readOptionalString();
248-
version = input.readString();
249-
protocol = input.readString();
256+
version = input.readOptionalString();
257+
protocol = input.readOptionalString();
250258
if (input.readBoolean()) {
251259
parameters = input.readMap(s -> s.readString(), s -> s.readString());
252260
}
@@ -268,5 +276,6 @@ public MLCreateConnectorInput(StreamInput input) throws IOException {
268276
this.access = input.readEnum(AccessMode.class);
269277
}
270278
dryRun = input.readBoolean();
279+
updateConnector = input.readBoolean();
271280
}
272281
}

‎common/src/main/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequest.java‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,32 +19,31 @@
1919
import java.io.ByteArrayOutputStream;
2020
import java.io.IOException;
2121
import java.io.UncheckedIOException;
22-
import java.util.Map;
2322

2423
import static org.opensearch.action.ValidateActions.addValidationError;
2524

2625
@Getter
2726
public class MLUpdateConnectorRequest extends ActionRequest {
2827
String connectorId;
29-
Map<String, Object> updateContent;
28+
MLCreateConnectorInput updateContent;
3029

3130
@Builder
32-
public MLUpdateConnectorRequest(String connectorId, Map<String, Object> updateContent) {
31+
public MLUpdateConnectorRequest(String connectorId, MLCreateConnectorInput updateContent) {
3332
this.connectorId = connectorId;
3433
this.updateContent = updateContent;
3534
}
3635

3736
public MLUpdateConnectorRequest(StreamInput in) throws IOException {
3837
super(in);
3938
this.connectorId = in.readString();
40-
this.updateContent = in.readMap();
39+
this.updateContent = new MLCreateConnectorInput(in);
4140
}
4241

4342
@Override
4443
public void writeTo(StreamOutput out) throws IOException {
4544
super.writeTo(out);
4645
out.writeString(this.connectorId);
47-
out.writeMap(this.getUpdateContent());
46+
this.updateContent.writeTo(out);
4847
}
4948

5049
@Override
@@ -55,14 +54,17 @@ public ActionRequestValidationException validate() {
5554
exception = addValidationError("ML connector id can't be null", exception);
5655
}
5756

57+
if (updateContent == null) {
58+
exception = addValidationError("Update connector content can't be null", exception);
59+
}
60+
5861
return exception;
5962
}
6063

6164
public static MLUpdateConnectorRequest parse(XContentParser parser, String connectorId) throws IOException {
62-
Map<String, Object> dataAsMap = null;
63-
dataAsMap = parser.map();
65+
MLCreateConnectorInput updateContent = MLCreateConnectorInput.parse(parser, true);
6466

65-
return MLUpdateConnectorRequest.builder().connectorId(connectorId).updateContent(dataAsMap).build();
67+
return MLUpdateConnectorRequest.builder().connectorId(connectorId).updateContent(updateContent).build();
6668
}
6769

6870
public static MLUpdateConnectorRequest fromActionRequest(ActionRequest actionRequest) {

‎common/src/test/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequestTests.java‎

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,38 +7,37 @@
77

88
import org.junit.Before;
99
import org.junit.Test;
10-
import org.mockito.Mock;
1110
import org.mockito.MockitoAnnotations;
1211
import org.opensearch.action.ActionRequest;
1312
import org.opensearch.action.ActionRequestValidationException;
1413
import org.opensearch.common.io.stream.BytesStreamOutput;
14+
import org.opensearch.common.settings.Settings;
15+
import org.opensearch.common.xcontent.XContentType;
1516
import org.opensearch.core.common.io.stream.StreamOutput;
17+
import org.opensearch.core.xcontent.NamedXContentRegistry;
1618
import org.opensearch.core.xcontent.XContentParser;
17-
import org.opensearch.rest.RestRequest;
19+
import org.opensearch.search.SearchModule;
1820

1921
import java.io.IOException;
2022
import java.io.UncheckedIOException;
21-
import java.util.Map;
23+
import java.util.Collections;
2224

2325
import static org.junit.Assert.assertEquals;
2426
import static org.junit.Assert.assertNotSame;
2527
import static org.junit.Assert.assertNull;
2628
import static org.junit.Assert.assertSame;
27-
import static org.mockito.Mockito.when;
29+
import static org.junit.Assert.assertTrue;
2830

2931
public class MLUpdateConnectorRequestTests {
3032
private String connectorId;
31-
private Map<String, Object> updateContent;
33+
private MLCreateConnectorInput updateContent;
3234
private MLUpdateConnectorRequest mlUpdateConnectorRequest;
3335

34-
@Mock
35-
XContentParser parser;
36-
3736
@Before
3837
public void setUp() {
3938
MockitoAnnotations.openMocks(this);
4039
this.connectorId = "test-connector_id";
41-
this.updateContent = Map.of("description", "new description");
40+
this.updateContent = MLCreateConnectorInput.builder().description("new description").updateConnector(true).build();
4241
mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
4342
.connectorId(connectorId)
4443
.updateContent(updateContent)
@@ -64,18 +63,20 @@ public void validate_Exception_NullConnectorId() {
6463
MLUpdateConnectorRequest updateConnectorRequest = MLUpdateConnectorRequest.builder().build();
6564
Exception exception = updateConnectorRequest.validate();
6665

67-
assertEquals("Validation Failed: 1: ML connector id can't be null;", exception.getMessage());
66+
assertEquals("Validation Failed: 1: ML connector id can't be null;2: Update connector content can't be null;", exception.getMessage());
6867
}
6968

7069
@Test
7170
public void parse_success() throws IOException {
72-
RestRequest.Method method = RestRequest.Method.POST;
73-
final Map<String, Object> updatefields = Map.of("version", "new version", "description", "new description");
74-
when(parser.map()).thenReturn(updatefields);
75-
71+
String jsonStr = "{\"version\":\"new version\",\"description\":\"new description\"}";
72+
XContentParser parser = XContentType.JSON.xContent().createParser(new NamedXContentRegistry(new SearchModule(Settings.EMPTY,
73+
Collections.emptyList()).getNamedXContents()), null, jsonStr);
74+
parser.nextToken();
7675
MLUpdateConnectorRequest updateConnectorRequest = MLUpdateConnectorRequest.parse(parser, connectorId);
7776
assertEquals(updateConnectorRequest.getConnectorId(), connectorId);
78-
assertEquals(updateConnectorRequest.getUpdateContent(), updatefields);
77+
assertTrue(updateConnectorRequest.getUpdateContent().isUpdateConnector());
78+
assertEquals("new version", updateConnectorRequest.getUpdateContent().getVersion());
79+
assertEquals("new description", updateConnectorRequest.getUpdateContent().getDescription());
7980
}
8081

8182
@Test

‎plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@
88
import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
99
import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
1010

11+
import java.util.ArrayList;
12+
import java.util.Arrays;
13+
import java.util.List;
14+
1115
import org.opensearch.action.ActionRequest;
1216
import org.opensearch.action.DocWriteResponse;
1317
import org.opensearch.action.delete.DeleteRequest;
@@ -77,11 +81,16 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
7781
.error(
7882
searchHits.length + " models are still using this connector, please delete or update the models first!"
7983
);
84+
List<String> modelIds = new ArrayList<>();
85+
for (SearchHit hit : searchHits) {
86+
modelIds.add(hit.getId());
87+
}
8088
actionListener
8189
.onFailure(
8290
new MLValidationException(
8391
searchHits.length
84-
+ " models are still using this connector, please delete or update the models first!"
92+
+ " models are still using this connector, please delete or update the models first: "
93+
+ Arrays.toString(modelIds.toArray(new String[0]))
8594
)
8695
);
8796
}

0 commit comments

Comments
 (0)