Get and delete agent APIs (opensearch-project#1752)

* get and delete agent APIs (opensearch-project#1703) Signed-off-by: Bhavana Ramaram <[email protected]> Signed-off-by: Mingshi Liu <[email protected]> * Add unit tests for Get and Delete APIs Signed-off-by: Mingshi Liu <[email protected]> * Add header and increase code coverage Signed-off-by: Mingshi Liu <[email protected]> * change IndexNotFoundException error message Signed-off-by: Mingshi Liu <[email protected]> --------- Signed-off-by: Bhavana Ramaram <[email protected]> Signed-off-by: Mingshi Liu <[email protected]> Signed-off-by: Mingshi Liu <[email protected]>
TrungBui59 · Dec 16, 2023 · 4d8d32e · 4d8d32e
1 parent bab9439
commit 4d8d32e
Show file tree

Hide file tree

Showing 20 changed files with 1,401 additions and 0 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteAction.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteAction.java
@@ -0,0 +1,16 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.action.delete.DeleteResponse;
+
+public class MLAgentDeleteAction extends ActionType<DeleteResponse> {
+    public static final MLAgentDeleteAction INSTANCE = new MLAgentDeleteAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/agents/delete";
+
+    private MLAgentDeleteAction() { super(NAME, DeleteResponse::new);}
+}
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteRequest.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteRequest.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+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 java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import static org.opensearch.action.ValidateActions.addValidationError;
+
+public class MLAgentDeleteRequest extends ActionRequest {
+    @Getter
+    String agentId;
+
+    @Builder
+    public MLAgentDeleteRequest(String agentId) {
+        this.agentId = agentId;
+    }
+
+    public MLAgentDeleteRequest(StreamInput input) throws IOException {
+        super(input);
+        this.agentId = input.readString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput output) throws IOException {
+        super.writeTo(output);
+        output.writeString(agentId);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException exception = null;
+
+        if (this.agentId == null) {
+            exception = addValidationError("ML agent id can't be null", exception);
+        }
+
+        return exception;
+    }
+
+    public static MLAgentDeleteRequest fromActionRequest(ActionRequest actionRequest) {
+        if (actionRequest instanceof MLAgentDeleteRequest) {
+            return (MLAgentDeleteRequest)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 MLAgentDeleteRequest(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionRequest into MLAgentDeleteRequest", e);
+        }
+    }
+
+}
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetAction.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetAction.java
@@ -0,0 +1,16 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import org.opensearch.action.ActionType;
+
+public class MLAgentGetAction extends ActionType<MLAgentGetResponse> {
+    public static final MLAgentGetAction INSTANCE = new MLAgentGetAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/agents/get";
+
+    private MLAgentGetAction() { super(NAME, MLAgentGetResponse::new);}
+
+}
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetRequest.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetRequest.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+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 java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import static org.opensearch.action.ValidateActions.addValidationError;
+
+@Getter
+public class MLAgentGetRequest extends ActionRequest {
+
+    String agentId;
+
+    @Builder
+    public MLAgentGetRequest(String agentId) {
+        this.agentId = agentId;
+    }
+
+    public MLAgentGetRequest(StreamInput in) throws IOException {
+        super(in);
+        this.agentId = in.readString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeString(this.agentId);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException exception = null;
+
+        if (this.agentId == null) {
+            exception = addValidationError("ML agent id can't be null", exception);
+        }
+
+        return exception;
+    }
+
+    public static MLAgentGetRequest fromActionRequest(ActionRequest actionRequest) {
+        if (actionRequest instanceof MLAgentGetRequest) {
+            return (MLAgentGetRequest) 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 MLAgentGetRequest(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionRequest into MLAgentGetRequest", e);
+        }
+    }
+}
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetResponse.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetResponse.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import lombok.Builder;
+import org.opensearch.core.action.ActionResponse;
+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.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.ml.common.agent.MLAgent;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+public class MLAgentGetResponse extends ActionResponse implements ToXContentObject {
+    MLAgent mlAgent;
+
+    @Builder
+    public MLAgentGetResponse(MLAgent mlAgent) {
+        this.mlAgent = mlAgent;
+    }
+
+    public MLAgentGetResponse(StreamInput in) throws IOException {
+        super(in);
+        mlAgent = MLAgent.fromStream(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException{
+        mlAgent.writeTo(out);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
+        return mlAgent.toXContent(xContentBuilder, params);
+    }
+
+    public static MLAgentGetResponse fromActionResponse(ActionResponse actionResponse) {
+        if (actionResponse instanceof MLAgentGetResponse) {
+            return (MLAgentGetResponse) actionResponse;
+        }
+
+        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
+             OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
+            actionResponse.writeTo(osso);
+            try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
+                return new MLAgentGetResponse(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionResponse into MLAgentGetResponse", e);
+        }
+    }
+
+}
diff --git a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteActionTest.java b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteActionTest.java
@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.ml.common.transport.agent;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+public class MLAgentDeleteActionTest {
+    @Test
+    public void testMLAgentDeleteActionInstance() {
+        assertNotNull(MLAgentDeleteAction.INSTANCE);
+        assertEquals("cluster:admin/opensearch/ml/agents/delete", MLAgentDeleteAction.NAME);
+    }
+
+}
diff --git a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteRequestTest.java b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteRequestTest.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.ml.common.transport.agent;
+
+import org.junit.Test;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+
+import java.io.IOException;
+
+import static org.junit.Assert.assertEquals;
+import static org.opensearch.action.ValidateActions.addValidationError;
+
+public class MLAgentDeleteRequestTest {
+    String agentId;
+
+    @Test
+    public void constructor_AgentId() {
+        agentId = "test-abc";
+        MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
+        assertEquals(mLAgentDeleteRequest.agentId,agentId);
+    }
+
+    @Test
+    public void writeTo() throws IOException {
+        agentId = "test-hij";
+
+        MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
+        BytesStreamOutput output = new BytesStreamOutput();
+        mLAgentDeleteRequest.writeTo(output);
+
+        MLAgentDeleteRequest mLAgentDeleteRequest1 = new MLAgentDeleteRequest(output.bytes().streamInput());
+
+        assertEquals(mLAgentDeleteRequest.agentId, mLAgentDeleteRequest1.agentId);
+        assertEquals(agentId, mLAgentDeleteRequest1.agentId);
+    }
+
+    @Test
+    public void validate_Success() {
+        agentId = "not-null";
+        MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
+
+        assertEquals(null, mLAgentDeleteRequest.validate());
+    }
+
+    @Test
+    public void validate_Failure() {
+        agentId = null;
+        MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
+        assertEquals(null,mLAgentDeleteRequest.agentId);
+
+        ActionRequestValidationException exception = addValidationError("ML agent id can't be null", null);
+        mLAgentDeleteRequest.validate().equals(exception) ;
+    }
+
+    @Test
+    public void fromActionRequest() throws IOException {
+        agentId = "test-lmn";
+        MLAgentDeleteRequest mLAgentDeleteRequest = new MLAgentDeleteRequest(agentId);
+        assertEquals(mLAgentDeleteRequest.fromActionRequest(mLAgentDeleteRequest), mLAgentDeleteRequest);
+
+    }
+}
diff --git a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentGetActionTest.java b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentGetActionTest.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import org.junit.Test;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+public class MLAgentGetActionTest {
+
+    @Test
+    public void testMLAgentGetActionInstance() {
+        assertNotNull(MLAgentGetAction.INSTANCE);
+        assertEquals("cluster:admin/opensearch/ml/agents/get", MLAgentGetAction.NAME);
+    }
+
+
+}