fix breaking changes in config index fields

Signed-off-by: Bhavana Ramaram <[email protected]>
opensearch-project · Sep 3, 2024 · 02b0d56 · 02b0d56
1 parent 49d4a01
commit 02b0d56
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 9 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/MLConfig.java b/common/src/main/java/org/opensearch/ml/common/MLConfig.java
@@ -42,55 +42,93 @@ public class MLConfig implements ToXContentObject, Writeable {
     @Setter
     private String type;
 
+    @Setter
+    private String configType;
+
     private Configuration configuration;
+    private Configuration mlConfiguration;
     private final Instant createTime;
     private Instant lastUpdateTime;
+    private Instant lastUpdatedTime;
 
     @Builder(toBuilder = true)
-    public MLConfig(String type, Configuration configuration, Instant createTime, Instant lastUpdateTime) {
+    public MLConfig(
+        String type,
+        String configType,
+        Configuration configuration,
+        Configuration mlConfiguration,
+        Instant createTime,
+        Instant lastUpdateTime,
+        Instant lastUpdatedTime
+    ) {
         this.type = type;
+        this.configType = configType;
         this.configuration = configuration;
+        this.mlConfiguration = mlConfiguration;
         this.createTime = createTime;
         this.lastUpdateTime = lastUpdateTime;
+        this.lastUpdatedTime = lastUpdatedTime;
     }
 
     public MLConfig(StreamInput input) throws IOException {
         this.type = input.readOptionalString();
+        this.configType = input.readOptionalString();
         if (input.readBoolean()) {
             configuration = new Configuration(input);
         }
+        if (input.readBoolean()) {
+            mlConfiguration = new Configuration(input);
+        }
         createTime = input.readOptionalInstant();
         lastUpdateTime = input.readOptionalInstant();
+        lastUpdatedTime = input.readOptionalInstant();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(type);
+        out.writeOptionalString(configType);
         if (configuration != null) {
             out.writeBoolean(true);
             configuration.writeTo(out);
         } else {
             out.writeBoolean(false);
         }
+        if (mlConfiguration != null) {
+            out.writeBoolean(true);
+            mlConfiguration.writeTo(out);
+        } else {
+            out.writeBoolean(false);
+        }
         out.writeOptionalInstant(createTime);
         out.writeOptionalInstant(lastUpdateTime);
+        out.writeOptionalInstant(lastUpdatedTime);
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
         XContentBuilder builder = xContentBuilder.startObject();
         if (type != null) {
-            builder.field(CONFIG_TYPE_FIELD, type);
+            builder.field(TYPE_FIELD, type);
+        }
+        if (configType != null) {
+            builder.field(CONFIG_TYPE_FIELD, configType);
         }
         if (configuration != null) {
-            builder.field(ML_CONFIGURATION_FIELD, configuration);
+            builder.field(CONFIGURATION_FIELD, configuration);
+        }
+        if (mlConfiguration != null) {
+            builder.field(ML_CONFIGURATION_FIELD, mlConfiguration);
         }
         if (createTime != null) {
             builder.field(CREATE_TIME_FIELD, createTime.toEpochMilli());
         }
         if (lastUpdateTime != null) {
             builder.field(LAST_UPDATE_TIME_FIELD, lastUpdateTime.toEpochMilli());
         }
+        if (lastUpdatedTime != null) {
+            builder.field(LAST_UPDATED_TIME_FIELD, lastUpdatedTime.toEpochMilli());
+        }
         return builder.endObject();
     }
 
@@ -142,10 +180,13 @@ public static MLConfig parse(XContentParser parser) throws IOException {
         }
         return MLConfig
             .builder()
-            .type(configType == null ? type : configType)
-            .configuration(mlConfiguration == null ? configuration : mlConfiguration)
+            .type(type)
+            .configType(configType)
+            .configuration(configuration)
+            .mlConfiguration(mlConfiguration)
             .createTime(createTime)
-            .lastUpdateTime(lastUpdatedTime == null ? lastUpdateTime : lastUpdatedTime)
+            .lastUpdateTime(lastUpdateTime)
+            .lastUpdatedTime(lastUpdatedTime)
             .build();
     }
 }
diff --git a/common/src/test/java/org/opensearch/ml/common/transport/config/MLConfigGetResponseTest.java b/common/src/test/java/org/opensearch/ml/common/transport/config/MLConfigGetResponseTest.java
@@ -60,7 +60,7 @@ public void MLConfigGetResponse_Builder() throws IOException {
     @Test
     public void writeTo() throws IOException {
         // create ml agent using mlConfig and mlConfigGetResponse
-        mlConfig = new MLConfig("olly_agent", new Configuration("agent_id"), Instant.EPOCH, Instant.EPOCH);
+        mlConfig = new MLConfig("olly_agent", null, new Configuration("agent_id"), null, Instant.EPOCH, Instant.EPOCH, Instant.EPOCH);
         MLConfigGetResponse mlConfigGetResponse = MLConfigGetResponse.builder().mlConfig(mlConfig).build();
         // use write out for both agents
         BytesStreamOutput output = new BytesStreamOutput();
@@ -76,7 +76,7 @@ public void writeTo() throws IOException {
 
     @Test
     public void toXContent() throws IOException {
-        mlConfig = new MLConfig(null, null, null, null);
+        mlConfig = new MLConfig(null, null, null, null, null, null, null);
         MLConfigGetResponse mlConfigGetResponse = MLConfigGetResponse.builder().mlConfig(mlConfig).build();
         XContentBuilder builder = XContentFactory.jsonBuilder();
         ToXContent.Params params = EMPTY_PARAMS;

diff --git a/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java
@@ -162,9 +162,32 @@ public void testDoExecute_Success() throws IOException {
         verify(actionListener).onResponse(any(MLConfigGetResponse.class));
     }
 
+    @Test
+    public void testDoExecute_Success_ForNewFields() throws IOException {
+        String configID = "config_id";
+        MLConfig mlConfig = new MLConfig(null, "olly_agent", null, new Configuration("agent_id"), Instant.EPOCH, null, Instant.EPOCH);
+
+        XContentBuilder content = mlConfig.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS);
+        BytesReference bytesReference = BytesReference.bytes(content);
+        GetResult getResult = new GetResult("indexName", configID, 111l, 111l, 111l, true, bytesReference, null, null);
+        GetResponse getResponse = new GetResponse(getResult);
+        ActionListener<MLConfigGetResponse> actionListener = mock(ActionListener.class);
+        MLConfigGetRequest request = new MLConfigGetRequest(configID);
+        Task task = mock(Task.class);
+
+        doAnswer(invocation -> {
+            ActionListener<GetResponse> listener = invocation.getArgument(1);
+            listener.onResponse(getResponse);
+            return null;
+        }).when(client).get(any(), any());
+
+        getConfigTransportAction.doExecute(task, request, actionListener);
+        verify(actionListener).onResponse(any(MLConfigGetResponse.class));
+    }
+
     public GetResponse prepareMLConfig(String configID) throws IOException {
 
-        MLConfig mlConfig = new MLConfig("olly_agent", new Configuration("agent_id"), Instant.EPOCH, Instant.EPOCH);
+        MLConfig mlConfig = new MLConfig("olly_agent", null, new Configuration("agent_id"), null, Instant.EPOCH, Instant.EPOCH, null);
 
         XContentBuilder content = mlConfig.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS);
         BytesReference bytesReference = BytesReference.bytes(content);