[Backport to main] update connector API (#1651)

* update connector API

Signed-off-by: Xun Zhang <[email protected]>

* more ut test coverage

Signed-off-by: Xun Zhang <[email protected]>

* check connector usage in deployed models before updating connector

Signed-off-by: Xun Zhang <[email protected]>

---------

Signed-off-by: Xun Zhang <[email protected]>
Co-authored-by: Xun Zhang <[email protected]>

common/src/main/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorAction.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.connector;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.action.update.UpdateResponse;
+
+public class MLUpdateConnectorAction extends ActionType<UpdateResponse> {
+    public static final MLUpdateConnectorAction INSTANCE = new MLUpdateConnectorAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/connectors/update";
+
+    private MLUpdateConnectorAction() { super(NAME, UpdateResponse::new);}
+}

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

@@ -0,0 +1,83 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.connector;
+
+import lombok.Builder;
+import lombok.Getter;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.core.common.io.stream.InputStreamStreamInput;
+import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.XContentParser;
+
+import java.io.ByteArrayInputStream;
+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;
+
+    @Builder
+    public MLUpdateConnectorRequest(String connectorId, Map<String, Object> updateContent) {
+        this.connectorId = connectorId;
+        this.updateContent = updateContent;
+    }
+
+    public MLUpdateConnectorRequest(StreamInput in) throws IOException {
+        super(in);
+        this.connectorId = in.readString();
+        this.updateContent = in.readMap();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeString(this.connectorId);
+        out.writeMap(this.getUpdateContent());
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException exception = null;
+
+        if (this.connectorId == null) {
+            exception = addValidationError("ML connector id 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();
+
+        return MLUpdateConnectorRequest.builder().connectorId(connectorId).updateContent(dataAsMap).build();
+    }
+
+    public static MLUpdateConnectorRequest fromActionRequest(ActionRequest actionRequest) {
+        if (actionRequest instanceof MLUpdateConnectorRequest) {
+            return (MLUpdateConnectorRequest) actionRequest;
+        }
+
+        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
+             OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
+            actionRequest.writeTo(osso);
+            try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
+                return new MLUpdateConnectorRequest(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionRequest into MLUpdateConnectorRequest", e);
+        }
+    }
+}

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

@@ -0,0 +1,128 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.connector;
+
+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.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.rest.RestRequest;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Map;
+
+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;
+
+public class MLUpdateConnectorRequestTests {
+    private String connectorId;
+    private Map<String, Object> 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");
+        mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
+            .connectorId(connectorId)
+            .updateContent(updateContent)
+            .build();
+    }
+
+    @Test
+    public void writeTo_Success() throws IOException {
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        mlUpdateConnectorRequest.writeTo(bytesStreamOutput);
+        MLUpdateConnectorRequest parsedUpdateRequest = new MLUpdateConnectorRequest(bytesStreamOutput.bytes().streamInput());
+        assertEquals(connectorId, parsedUpdateRequest.getConnectorId());
+        assertEquals(updateContent, parsedUpdateRequest.getUpdateContent());
+    }
+
+    @Test
+    public void validate_Success() {
+        assertNull(mlUpdateConnectorRequest.validate());
+    }
+
+    @Test
+    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());
+    }
+
+    @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);
+
+        MLUpdateConnectorRequest updateConnectorRequest = MLUpdateConnectorRequest.parse(parser, connectorId);
+        assertEquals(updateConnectorRequest.getConnectorId(), connectorId);
+        assertEquals(updateConnectorRequest.getUpdateContent(), updatefields);
+    }
+
+    @Test
+    public void fromActionRequest_Success() {
+        MLUpdateConnectorRequest mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
+            .connectorId(connectorId)
+            .updateContent(updateContent)
+            .build();
+        assertSame(MLUpdateConnectorRequest.fromActionRequest(mlUpdateConnectorRequest), mlUpdateConnectorRequest);
+    }
+
+    @Test
+    public void fromActionRequest_Success_fromActionRequest() {
+        MLUpdateConnectorRequest mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
+            .connectorId(connectorId)
+            .updateContent(updateContent)
+            .build();
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                mlUpdateConnectorRequest.writeTo(out);
+            }
+        };
+        MLUpdateConnectorRequest request = MLUpdateConnectorRequest.fromActionRequest(actionRequest);
+        assertNotSame(request, mlUpdateConnectorRequest);
+        assertEquals(mlUpdateConnectorRequest.getConnectorId(), request.getConnectorId());
+        assertEquals(mlUpdateConnectorRequest.getUpdateContent(), request.getUpdateContent());
+    }
+
+    @Test(expected = UncheckedIOException.class)
+    public void fromActionRequest_IOException() {
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                throw new IOException();
+            }
+        };
+        MLUpdateConnectorRequest.fromActionRequest(actionRequest);
+    }
+}

plugin/src/main/java/org/opensearch/ml/action/connector/UpdateConnectorTransportAction.java

@@ -0,0 +1,146 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.action.connector;
+
+import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
+
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.DocWriteResponse;
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.action.update.UpdateRequest;
+import org.opensearch.action.update.UpdateResponse;
+import org.opensearch.client.Client;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.index.IndexNotFoundException;
+import org.opensearch.index.query.BoolQueryBuilder;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.ml.common.MLModel;
+import org.opensearch.ml.common.exception.MLValidationException;
+import org.opensearch.ml.common.transport.connector.MLUpdateConnectorAction;
+import org.opensearch.ml.common.transport.connector.MLUpdateConnectorRequest;
+import org.opensearch.ml.helper.ConnectorAccessControlHelper;
+import org.opensearch.ml.model.MLModelManager;
+import org.opensearch.search.SearchHit;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.tasks.Task;
+import org.opensearch.transport.TransportService;
+
+import lombok.AccessLevel;
+import lombok.experimental.FieldDefaults;
+import lombok.extern.log4j.Log4j2;
+
+@Log4j2
+@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
+public class UpdateConnectorTransportAction extends HandledTransportAction<ActionRequest, UpdateResponse> {
+    Client client;
+
+    ConnectorAccessControlHelper connectorAccessControlHelper;
+    MLModelManager mlModelManager;
+
+    @Inject
+    public UpdateConnectorTransportAction(
+        TransportService transportService,
+        ActionFilters actionFilters,
+        Client client,
+        ConnectorAccessControlHelper connectorAccessControlHelper,
+        MLModelManager mlModelManager
+    ) {
+        super(MLUpdateConnectorAction.NAME, transportService, actionFilters, MLUpdateConnectorRequest::new);
+        this.client = client;
+        this.connectorAccessControlHelper = connectorAccessControlHelper;
+        this.mlModelManager = mlModelManager;
+    }
+
+    @Override
+    protected void doExecute(Task task, ActionRequest request, ActionListener<UpdateResponse> listener) {
+        MLUpdateConnectorRequest mlUpdateConnectorAction = MLUpdateConnectorRequest.fromActionRequest(request);
+        String connectorId = mlUpdateConnectorAction.getConnectorId();
+        UpdateRequest updateRequest = new UpdateRequest(ML_CONNECTOR_INDEX, connectorId);
+        updateRequest.doc(mlUpdateConnectorAction.getUpdateContent());
+        updateRequest.docAsUpsert(true);
+
+        try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+            connectorAccessControlHelper.validateConnectorAccess(client, connectorId, ActionListener.wrap(hasPermission -> {
+                if (Boolean.TRUE.equals(hasPermission)) {
+                    updateUndeployedConnector(connectorId, updateRequest, listener, context);
+                } else {
+                    listener
+                        .onFailure(
+                            new IllegalArgumentException("You don't have permission to update the connector, connector id: " + connectorId)
+                        );
+                }
+            }, exception -> {
+                log.error("Permission denied: Unable to update the connector with ID {}. Details: {}", connectorId, exception);
+                listener.onFailure(exception);
+            }));
+        } catch (Exception e) {
+            log.error("Failed to update ML connector for connector id {}. Details {}:", connectorId, e);
+            listener.onFailure(e);
+        }
+    }
+
+    private void updateUndeployedConnector(
+        String connectorId,
+        UpdateRequest updateRequest,
+        ActionListener<UpdateResponse> listener,
+        ThreadContext.StoredContext context
+    ) {
+        SearchRequest searchRequest = new SearchRequest(ML_MODEL_INDEX);
+        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
+        boolQueryBuilder.must(QueryBuilders.matchQuery(MLModel.CONNECTOR_ID_FIELD, connectorId));
+        boolQueryBuilder.must(QueryBuilders.idsQuery().addIds(mlModelManager.getAllModelIds()));
+        sourceBuilder.query(boolQueryBuilder);
+        searchRequest.source(sourceBuilder);
+
+        client.search(searchRequest, ActionListener.wrap(searchResponse -> {
+            SearchHit[] searchHits = searchResponse.getHits().getHits();
+            if (searchHits.length == 0) {
+                client.update(updateRequest, getUpdateResponseListener(connectorId, listener, context));
+            } else {
+                log.error(searchHits.length + " models are still using this connector, please undeploy the models first!");
+                listener
+                    .onFailure(
+                        new MLValidationException(
+                            searchHits.length + " models are still using this connector, please undeploy the models first!"
+                        )
+                    );
+            }
+        }, e -> {
+            if (e instanceof IndexNotFoundException) {
+                client.update(updateRequest, getUpdateResponseListener(connectorId, listener, context));
+                return;
+            }
+            log.error("Failed to update ML connector: " + connectorId, e);
+            listener.onFailure(e);
+
+        }));
+    }
+
+    private ActionListener<UpdateResponse> getUpdateResponseListener(
+        String connectorId,
+        ActionListener<UpdateResponse> actionListener,
+        ThreadContext.StoredContext context
+    ) {
+        return ActionListener.runBefore(ActionListener.wrap(updateResponse -> {
+            if (updateResponse != null && updateResponse.getResult() != DocWriteResponse.Result.UPDATED) {
+                log.info("Failed to update the connector with ID: {}", connectorId);
+                actionListener.onResponse(updateResponse);
+                return;
+            }
+            log.info("Successfully updated the connector with ID: {}", connectorId);
+            actionListener.onResponse(updateResponse);
+        }, exception -> {
+            log.error("Failed to update ML connector with ID {}. Details: {}", connectorId, exception);
+            actionListener.onFailure(exception);
+        }), context::restore);
+    }
+}