From 6b925be86042a81f7fe12376f440f5712c443876 Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Tue, 16 Jul 2024 15:38:03 -0700 Subject: [PATCH] Fix SearchSourceBuilder read/write test failures Signed-off-by: Siddhant Deshmukh --- .../opensearch/plugin/insights/rules/model/Attribute.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java b/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java index 9d9f3db5..43d7f61c 100644 --- a/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java +++ b/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java @@ -13,6 +13,7 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo; +import org.opensearch.search.builder.SearchSourceBuilder; import java.io.IOException; import java.util.Collections; @@ -90,6 +91,8 @@ static void writeTo(final StreamOutput out, final Attribute attribute) throws IO public static void writeValueTo(StreamOutput out, Object attributeValue) throws IOException { if (attributeValue instanceof List) { out.writeList((List) attributeValue); + } else if (attributeValue instanceof SearchSourceBuilder) { + ((SearchSourceBuilder) attributeValue).writeTo(out); } else { out.writeGenericValue(attributeValue); } @@ -106,6 +109,9 @@ public static void writeValueTo(StreamOutput out, Object attributeValue) throws public static Object readAttributeValue(StreamInput in, Attribute attribute) throws IOException { if (attribute == Attribute.TASK_RESOURCE_USAGES) { return in.readList(TaskResourceInfo::readFromStream); + } else if (attribute == Attribute.SOURCE) { + SearchSourceBuilder builder = new SearchSourceBuilder(in); + return builder; } else { return in.readGenericValue(); }