fix update connector API (#1484) (#1494)

* fix update connector API Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix ut Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 2f3f39e) Co-authored-by: Yaliang Wu <ylwu@amazon.com>
opensearch-project · Oct 11, 2023 · 61012a5 · 61012a5
1 parent cea1cd6
commit 61012a5
Show file tree

Hide file tree

Showing 13 changed files with 275 additions and 133 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/connector/Connector.java b/common/src/main/java/org/opensearch/ml/common/connector/Connector.java
@@ -30,17 +30,7 @@
 import org.opensearch.ml.common.AccessMode;
 import org.opensearch.ml.common.MLCommonsClassLoader;
 import org.opensearch.ml.common.output.model.ModelTensor;
-
-import java.io.IOException;
-import java.security.AccessController;
-import java.security.PrivilegedActionException;
-import java.security.PrivilegedExceptionAction;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.function.Function;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
+import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.ml.common.utils.StringUtils.gson;
@@ -80,6 +70,7 @@ public interface Connector extends ToXContentObject, Writeable {
 
     void writeTo(StreamOutput out) throws IOException;
 
+    void update(MLCreateConnectorInput updateContent, Function<String, String> function);
 
     <T> void parseResponse(T orElse, List<ModelTensor> modelTensors, boolean b) throws IOException;
 

diff --git a/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java b/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java
@@ -34,6 +34,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.ml.common.AccessMode;
+import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
 
 @Log4j2
 @NoArgsConstructor
@@ -248,6 +249,38 @@ public void writeTo(StreamOutput out) throws IOException {
         }
     }
 
+    @Override
+    public void update(MLCreateConnectorInput updateContent, Function<String, String> function) {
+        if (updateContent.getName() != null) {
+            this.name = updateContent.getName();
+        }
+        if (updateContent.getDescription() != null) {
+            this.description = updateContent.getDescription();
+        }
+        if (updateContent.getVersion() != null) {
+            this.version = updateContent.getVersion();
+        }
+        if (updateContent.getProtocol() != null) {
+            this.protocol = updateContent.getProtocol();
+        }
+        if (updateContent.getParameters() != null && updateContent.getParameters().size() > 0) {
+            this.parameters = updateContent.getParameters();
+        }
+        if (updateContent.getCredential() != null && updateContent.getCredential().size() > 0) {
+            this.credential = updateContent.getCredential();
+            encrypt(function);
+        }
+        if (updateContent.getActions() != null) {
+            this.actions = updateContent.getActions();
+        }
+        if (updateContent.getBackendRoles() != null) {
+            this.backendRoles = updateContent.getBackendRoles();
+        }
+        if (updateContent.getAccess() != null) {
+            this.access = updateContent.getAccess();
+        }
+    }
+
     @Override
     public  <T> T createPredictPayload(Map<String, String> parameters) {
         Optional<ConnectorAction> predictAction = findPredictAction();

diff --git a/...on/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java b/...on/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java
@@ -56,6 +56,7 @@ public class MLCreateConnectorInput implements ToXContentObject, Writeable {
     private Boolean addAllBackendRoles;
     private AccessMode access;
     private boolean dryRun = false;
+    private boolean updateConnector = false;
 
     @Builder(toBuilder = true)
     public MLCreateConnectorInput(String name,
@@ -68,9 +69,10 @@ public MLCreateConnectorInput(String name,
                                   List<String> backendRoles,
                                   Boolean addAllBackendRoles,
                                   AccessMode access,
-                                  boolean dryRun
+                                  boolean dryRun,
+                                  boolean updateConnector
     ) {
-        if (!dryRun) {
+        if (!dryRun && !updateConnector) {
             if (name == null) {
                 throw new IllegalArgumentException("Connector name is null");
             }
@@ -92,9 +94,14 @@ public MLCreateConnectorInput(String name,
         this.addAllBackendRoles = addAllBackendRoles;
         this.access = access;
         this.dryRun = dryRun;
+        this.updateConnector = updateConnector;
     }
 
     public static MLCreateConnectorInput parse(XContentParser parser) throws IOException {
+        return parse(parser, false);
+    }
+
+    public static MLCreateConnectorInput parse(XContentParser parser, boolean updateConnector) throws IOException {
         String name = null;
         String description = null;
         String version = null;
@@ -159,7 +166,7 @@ public static MLCreateConnectorInput parse(XContentParser parser) throws IOExcep
                     break;
             }
         }
-        return new MLCreateConnectorInput(name, description, version, protocol, parameters, credential, actions, backendRoles, addAllBackendRoles, access, dryRun);
+        return new MLCreateConnectorInput(name, description, version, protocol, parameters, credential, actions, backendRoles, addAllBackendRoles, access, dryRun, updateConnector);
     }
 
     @Override
@@ -201,10 +208,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
     @Override
     public void writeTo(StreamOutput output) throws IOException {
-        output.writeString(name);
+        output.writeOptionalString(name);
         output.writeOptionalString(description);
-        output.writeString(version);
-        output.writeString(protocol);
+        output.writeOptionalString(version);
+        output.writeOptionalString(protocol);
         if (parameters != null) {
             output.writeBoolean(true);
             output.writeMap(parameters, StreamOutput::writeString, StreamOutput::writeString);
@@ -240,13 +247,14 @@ public void writeTo(StreamOutput output) throws IOException {
             output.writeBoolean(false);
         }
         output.writeBoolean(dryRun);
+        output.writeBoolean(updateConnector);
     }
 
     public MLCreateConnectorInput(StreamInput input) throws IOException {
-        name = input.readString();
+        name = input.readOptionalString();
         description = input.readOptionalString();
-        version = input.readString();
-        protocol = input.readString();
+        version = input.readOptionalString();
+        protocol = input.readOptionalString();
         if (input.readBoolean()) {
             parameters = input.readMap(s -> s.readString(), s -> s.readString());
         }
@@ -268,5 +276,6 @@ public MLCreateConnectorInput(StreamInput input) throws IOException {
             this.access = input.readEnum(AccessMode.class);
         }
         dryRun = input.readBoolean();
+        updateConnector = input.readBoolean();
     }
 }
diff --git a/.../src/main/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequest.java b/.../src/main/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequest.java
@@ -19,32 +19,31 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
-import java.util.Map;
 
 import static org.opensearch.action.ValidateActions.addValidationError;
 
 @Getter
 public class MLUpdateConnectorRequest extends ActionRequest {
     String connectorId;
-    Map<String, Object> updateContent;
+    MLCreateConnectorInput updateContent;
 
     @Builder
-    public MLUpdateConnectorRequest(String connectorId, Map<String, Object> updateContent) {
+    public MLUpdateConnectorRequest(String connectorId, MLCreateConnectorInput updateContent) {
         this.connectorId = connectorId;
         this.updateContent = updateContent;
     }
 
     public MLUpdateConnectorRequest(StreamInput in) throws IOException {
         super(in);
         this.connectorId = in.readString();
-        this.updateContent = in.readMap();
+        this.updateContent = new MLCreateConnectorInput(in);
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeString(this.connectorId);
-        out.writeMap(this.getUpdateContent());
+        this.updateContent.writeTo(out);
     }
 
     @Override
@@ -55,14 +54,17 @@ public ActionRequestValidationException validate() {
             exception = addValidationError("ML connector id can't be null", exception);
         }
 
+        if (updateContent == null) {
+            exception = addValidationError("Update connector content can't be null", exception);
+        }
+
         return exception;
     }
 
     public static MLUpdateConnectorRequest parse(XContentParser parser, String connectorId) throws IOException {
-        Map<String, Object> dataAsMap = null;
-        dataAsMap = parser.map();
+        MLCreateConnectorInput updateContent = MLCreateConnectorInput.parse(parser, true);
 
-        return MLUpdateConnectorRequest.builder().connectorId(connectorId).updateContent(dataAsMap).build();
+        return MLUpdateConnectorRequest.builder().connectorId(connectorId).updateContent(updateContent).build();
     }
 
     public static MLUpdateConnectorRequest fromActionRequest(ActionRequest actionRequest) {

diff --git a/...test/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequestTests.java b/...test/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequestTests.java
@@ -7,38 +7,37 @@
 
 import org.junit.Before;
 import org.junit.Test;
-import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.rest.RestRequest;
+import org.opensearch.search.SearchModule;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
-import java.util.Map;
+import java.util.Collections;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
-import static org.mockito.Mockito.when;
+import static org.junit.Assert.assertTrue;
 
 public class MLUpdateConnectorRequestTests {
     private String connectorId;
-    private Map<String, Object> updateContent;
+    private MLCreateConnectorInput updateContent;
     private MLUpdateConnectorRequest mlUpdateConnectorRequest;
 
-    @Mock
-    XContentParser parser;
-
     @Before
     public void setUp() {
         MockitoAnnotations.openMocks(this);
         this.connectorId = "test-connector_id";
-        this.updateContent = Map.of("description", "new description");
+        this.updateContent = MLCreateConnectorInput.builder().description("new description").updateConnector(true).build();
         mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
             .connectorId(connectorId)
             .updateContent(updateContent)
@@ -64,18 +63,20 @@ public void validate_Exception_NullConnectorId() {
         MLUpdateConnectorRequest updateConnectorRequest = MLUpdateConnectorRequest.builder().build();
         Exception exception = updateConnectorRequest.validate();
 
-        assertEquals("Validation Failed: 1: ML connector id can't be null;", exception.getMessage());
+        assertEquals("Validation Failed: 1: ML connector id can't be null;2: Update connector content can't be null;", exception.getMessage());
     }
 
     @Test
     public void parse_success() throws IOException {
-        RestRequest.Method method = RestRequest.Method.POST;
-        final Map<String, Object> updatefields = Map.of("version", "new version", "description", "new description");
-        when(parser.map()).thenReturn(updatefields);
-
+        String jsonStr = "{\"version\":\"new version\",\"description\":\"new description\"}";
+        XContentParser parser = XContentType.JSON.xContent().createParser(new NamedXContentRegistry(new SearchModule(Settings.EMPTY,
+                Collections.emptyList()).getNamedXContents()), null, jsonStr);
+        parser.nextToken();
         MLUpdateConnectorRequest updateConnectorRequest = MLUpdateConnectorRequest.parse(parser, connectorId);
         assertEquals(updateConnectorRequest.getConnectorId(), connectorId);
-        assertEquals(updateConnectorRequest.getUpdateContent(), updatefields);
+        assertTrue(updateConnectorRequest.getUpdateContent().isUpdateConnector());
+        assertEquals("new version", updateConnectorRequest.getUpdateContent().getVersion());
+        assertEquals("new description", updateConnectorRequest.getUpdateContent().getDescription());
     }
 
     @Test

diff --git a/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java
@@ -8,6 +8,10 @@
 import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.DocWriteResponse;
 import org.opensearch.action.delete.DeleteRequest;
@@ -77,11 +81,16 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                                 .error(
                                     searchHits.length + " models are still using this connector, please delete or update the models first!"
                                 );
+                            List<String> modelIds = new ArrayList<>();
+                            for (SearchHit hit : searchHits) {
+                                modelIds.add(hit.getId());
+                            }
                             actionListener
                                 .onFailure(
                                     new MLValidationException(
                                         searchHits.length
-                                            + " models are still using this connector, please delete or update the models first!"
+                                            + " models are still using this connector, please delete or update the models first: "
+                                            + Arrays.toString(modelIds.toArray(new String[0]))
                                     )
                                 );
                         }