add get config api to retrieve root agent id (#1995) (#1996)

add get config api to retrieve root agent id

Signed-off-by: Bhavana Ramaram <[email protected]>
(cherry picked from commit 9427f22)

Co-authored-by: Bhavana Ramaram <[email protected]>

common/src/main/java/org/opensearch/ml/common/CommonValue.java

@@ -45,6 +45,7 @@ public class CommonValue {
 
         public static final String MASTER_KEY = "master_key";
         public static final String CREATE_TIME_FIELD = "create_time";
+        public static final String LAST_UPDATE_TIME_FIELD = "last_update_time";
 
         public static final String BOX_TYPE_KEY = "box_type";
         // hot node
@@ -359,7 +360,16 @@ public class CommonValue {
                         + MASTER_KEY
                         + "\": {\"type\": \"keyword\"},\n"
                         + "      \""
+                        + MLConfig.TYPE_FIELD
+                        + "\" : {\"type\":\"keyword\"},\n"
+                        + "      \""
+                        + MLConfig.CONFIGURATION_FIELD
+                        + "\" : {\"type\": \"flat_object\"},\n"
+                        + "      \""
                         + CREATE_TIME_FIELD
+                        + "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
+                        + "      \""
+                        + LAST_UPDATE_TIME_FIELD
                         + "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"}\n"
                         + "    }\n"
                         + "}";

common/src/main/java/org/opensearch/ml/common/Configuration.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common;
+
+import lombok.Builder;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.xcontent.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParser;
+
+import java.io.IOException;
+
+import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+
+@Getter
+@EqualsAndHashCode
+public class Configuration implements ToXContentObject, Writeable {
+
+    public static final String ROOT_AGENT_ID = "agent_id";
+
+    @Setter
+    private String agentId;
+
+    @Builder(toBuilder = true)
+    public Configuration(
+        String agentId
+    ) {
+        this.agentId = agentId;
+    }
+
+    public Configuration(StreamInput input) throws IOException {
+        this.agentId = input.readOptionalString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeOptionalString(agentId);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
+        XContentBuilder builder = xContentBuilder.startObject();
+        if (agentId != null) {
+            builder.field(ROOT_AGENT_ID, agentId);
+        }
+        return builder.endObject();
+    }
+
+    public static Configuration fromStream(StreamInput in) throws IOException {
+        Configuration configuration = new Configuration(in);
+        return configuration;
+    }
+
+    public static Configuration parse(XContentParser parser) throws IOException {
+        String agentId = null;
+
+        ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
+        while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
+            String fieldName = parser.currentName();
+            parser.nextToken();
+
+            switch (fieldName) {
+                case ROOT_AGENT_ID:
+                    agentId = parser.text();
+                    break;
+                default:
+                    parser.skipChildren();
+                    break;
+            }
+        }
+        return Configuration.builder()
+                .agentId(agentId)
+                .build();
+    }
+}

common/src/main/java/org/opensearch/ml/common/MLConfig.java

@@ -0,0 +1,136 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common;
+
+import lombok.Builder;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.xcontent.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.time.Instant;
+
+import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+
+@Getter
+@EqualsAndHashCode
+public class MLConfig implements ToXContentObject, Writeable {
+
+    public static final String TYPE_FIELD = "type";
+
+    public static final String CONFIGURATION_FIELD = "configuration";
+
+    public static final String CREATE_TIME_FIELD = "create_time";
+    public static final String LAST_UPDATE_TIME_FIELD = "last_update_time";
+
+    @Setter
+    private String type;
+
+    private Configuration configuration;
+    private final Instant createTime;
+    private Instant lastUpdateTime;
+
+    @Builder(toBuilder = true)
+    public MLConfig(
+            String type,
+            Configuration configuration,
+            Instant createTime,
+            Instant lastUpdateTime
+    ) {
+        this.type = type;
+        this.configuration = configuration;
+        this.createTime = createTime;
+        this.lastUpdateTime = lastUpdateTime;
+    }
+
+    public MLConfig(StreamInput input) throws IOException {
+        this.type = input.readOptionalString();
+        if (input.readBoolean()) {
+            configuration = new Configuration(input);
+        }
+        createTime = input.readOptionalInstant();
+        lastUpdateTime = input.readOptionalInstant();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeOptionalString(type);
+        if (configuration != null) {
+            out.writeBoolean(true);
+            configuration.writeTo(out);
+        } else {
+            out.writeBoolean(false);
+        }
+        out.writeOptionalInstant(createTime);
+        out.writeOptionalInstant(lastUpdateTime);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
+        XContentBuilder builder = xContentBuilder.startObject();
+        if (type != null) {
+            builder.field(TYPE_FIELD, type);
+        }
+        if (configuration != null) {
+            builder.field(CONFIGURATION_FIELD, configuration);
+        }
+        if (createTime != null) {
+            builder.field(CREATE_TIME_FIELD, createTime.toEpochMilli());
+        }
+        if (lastUpdateTime != null) {
+            builder.field(LAST_UPDATE_TIME_FIELD, lastUpdateTime.toEpochMilli());
+        }
+        return builder.endObject();
+    }
+
+    public static MLConfig fromStream(StreamInput in) throws IOException {
+        MLConfig mlConfig = new MLConfig(in);
+        return mlConfig;
+    }
+
+    public static MLConfig parse(XContentParser parser) throws IOException {
+        String type = null;
+        Configuration configuration = null;
+        Instant createTime = null;
+        Instant lastUpdateTime = null;
+
+        ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
+        while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
+            String fieldName = parser.currentName();
+            parser.nextToken();
+
+            switch (fieldName) {
+                case TYPE_FIELD:
+                    type = parser.text();
+                    break;
+                case CONFIGURATION_FIELD:
+                    configuration = Configuration.parse(parser);
+                    break;
+                case CREATE_TIME_FIELD:
+                    createTime = Instant.ofEpochMilli(parser.longValue());
+                    break;
+                case LAST_UPDATE_TIME_FIELD:
+                    lastUpdateTime = Instant.ofEpochMilli(parser.longValue());
+                    break;
+                default:
+                    parser.skipChildren();
+                    break;
+            }
+        }
+        return MLConfig.builder()
+                .type(type)
+                .configuration(configuration)
+                .createTime(createTime)
+                .lastUpdateTime(lastUpdateTime)
+                .build();
+    }
+}

common/src/main/java/org/opensearch/ml/common/transport/config/MLConfigGetAction.java

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

common/src/main/java/org/opensearch/ml/common/transport/config/MLConfigGetRequest.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.config;
+
+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 MLConfigGetRequest extends ActionRequest {
+
+    String configId;
+
+    @Builder
+    public MLConfigGetRequest(String configId) {
+        this.configId = configId;
+    }
+
+    public MLConfigGetRequest(StreamInput in) throws IOException {
+        super(in);
+        this.configId = in.readString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeString(this.configId);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException exception = null;
+
+        if (this.configId == null) {
+            exception = addValidationError("ML config id can't be null", exception);
+        }
+
+        return exception;
+    }
+
+    public static MLConfigGetRequest fromActionRequest(ActionRequest actionRequest) {
+        if (actionRequest instanceof MLConfigGetRequest) {
+            return (MLConfigGetRequest) 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 MLConfigGetRequest(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionRequest into MLConfigGetRequest", e);
+        }
+    }
+}

common/src/main/java/org/opensearch/ml/common/transport/config/MLConfigGetResponse.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.config;
+
+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.MLConfig;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+public class MLConfigGetResponse extends ActionResponse implements ToXContentObject {
+    MLConfig mlConfig;
+
+    @Builder
+    public MLConfigGetResponse(MLConfig mlConfig) {
+        this.mlConfig = mlConfig;
+    }
+
+    public MLConfigGetResponse(StreamInput in) throws IOException {
+        super(in);
+        mlConfig = MLConfig.fromStream(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException{
+        mlConfig.writeTo(out);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
+        return mlConfig.toXContent(xContentBuilder, params);
+    }
+
+    public static MLConfigGetResponse fromActionResponse(ActionResponse actionResponse) {
+        if (actionResponse instanceof MLConfigGetResponse) {
+            return (MLConfigGetResponse) 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 MLConfigGetResponse(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionResponse into MLConfigGetResponse", e);
+        }
+    }
+
+}