From 42550cb6b9813e4a29fbf77f898ad1398b906cfe Mon Sep 17 00:00:00 2001
From: Thomas Farr <tsfarr@amazon.com>
Date: Thu, 10 Oct 2024 08:49:20 +1300
Subject: [PATCH] Fix MultisearchBody as well

Signed-off-by: Thomas Farr <tsfarr@amazon.com>
---
 CHANGELOG.md                                  |  1 +
 .../core/msearch/MultisearchBody.java         | 53 ++++++++++++++++---
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2b42b667de..1ba8e7dc07 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ### Added
 - Added `hashCode` and `equals` implementations ([#312](https://github.com/opensearch-project/opensearch-java/pull/312))
 - Added `sortVals` to `Hit` to allow retrieving non-string sort values ([#1224](https://github.com/opensearch-project/opensearch-java/pull/1224))
+- Added `searchAfterVals` to `MultisearchBody` to allow passing arbitrary `FieldValue`s to `search_after` ([#1224](https://github.com/opensearch-project/opensearch-java/pull/1224))
 
 ### Dependencies
 - Bumps `org.junit:junit-bom` from 5.10.3 to 5.11.1
diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/core/msearch/MultisearchBody.java b/java-client/src/main/java/org/opensearch/client/opensearch/core/msearch/MultisearchBody.java
index fb4ff1013d..4fc94119d5 100644
--- a/java-client/src/main/java/org/opensearch/client/opensearch/core/msearch/MultisearchBody.java
+++ b/java-client/src/main/java/org/opensearch/client/opensearch/core/msearch/MultisearchBody.java
@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import org.opensearch.client.json.JsonData;
 import org.opensearch.client.json.JsonpDeserializable;
@@ -44,6 +45,7 @@
 import org.opensearch.client.json.ObjectBuilderDeserializer;
 import org.opensearch.client.json.ObjectDeserializer;
 import org.opensearch.client.json.PlainJsonSerializable;
+import org.opensearch.client.opensearch._types.FieldValue;
 import org.opensearch.client.opensearch._types.ScriptField;
 import org.opensearch.client.opensearch._types.SortOptions;
 import org.opensearch.client.opensearch._types.aggregations.Aggregation;
@@ -77,7 +79,7 @@ public class MultisearchBody implements PlainJsonSerializable {
     @Nullable
     private final Query postFilter;
 
-    private final List<String> searchAfter;
+    private final List<FieldValue> searchAfter;
 
     @Nullable
     private final Integer size;
@@ -201,8 +203,19 @@ public final Query postFilter() {
 
     /**
      * API name: {@code search_after}
+     *
+     * <p><b>NOTE: In version 3.0.0 of opensearch-java, this method will instead return a {@code List<FieldValue>}.</b></p>
      */
     public final List<String> searchAfter() {
+        return this.searchAfter.stream().map(FieldValue::_toJsonString).collect(Collectors.toList());
+    }
+
+    /**
+     * API name: {@code search_after}
+     *
+     * <p><b>NOTE: In version 3.0.0 of opensearch-java, this method will be renamed to replace {@link #searchAfter()}.</b></p>
+     */
+    public final List<FieldValue> searchAfterVals() {
         return this.searchAfter;
     }
 
@@ -407,8 +420,8 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
         if (ApiTypeHelper.isDefined(this.searchAfter)) {
             generator.writeKey("search_after");
             generator.writeStartArray();
-            for (String item0 : this.searchAfter) {
-                generator.write(item0);
+            for (FieldValue item0 : this.searchAfter) {
+                item0.serialize(generator, mapper);
 
             }
             generator.writeEnd();
@@ -578,7 +591,7 @@ public static class Builder extends ObjectBuilderBase implements ObjectBuilder<M
         private Query postFilter;
 
         @Nullable
-        private List<String> searchAfter;
+        private List<FieldValue> searchAfter;
 
         @Nullable
         private Integer size;
@@ -715,9 +728,11 @@ public final Builder postFilter(Function<Query.Builder, ObjectBuilder<Query>> fn
          * API name: {@code search_after}
          * <p>
          * Adds all elements of <code>list</code> to <code>searchAfter</code>.
+         *
+         * <p><b>NOTE: In version 3.0.0 of opensearch-java, this method will instead accept a {@code List<FieldValue>}.</b></p>
          */
         public final Builder searchAfter(List<String> list) {
-            this.searchAfter = _listAddAll(this.searchAfter, list);
+            this.searchAfter = _listAddAll(this.searchAfter, FieldValue::of, list);
             return this;
         }
 
@@ -725,8 +740,34 @@ public final Builder searchAfter(List<String> list) {
          * API name: {@code search_after}
          * <p>
          * Adds one or more values to <code>searchAfter</code>.
+         *
+         * <p><b>NOTE: In version 3.0.0 of opensearch-java, this method will instead accept values of type {@code FieldValue}.</b></p>
          */
         public final Builder searchAfter(String value, String... values) {
+            this.searchAfter = _listAdd(this.searchAfter, FieldValue::of, value, values);
+            return this;
+        }
+
+        /**
+         * API name: {@code search_after}
+         * <p>
+         * Adds all elements of <code>list</code> to <code>searchAfter</code>.
+         *
+         * <p><b>NOTE: In version 3.0.0 of opensearch-java, this method will be renamed to replace {@link #searchAfter(List)}.</b></p>
+         */
+        public final Builder searchAfterVals(List<FieldValue> list) {
+            this.searchAfter = _listAddAll(this.searchAfter, list);
+            return this;
+        }
+
+        /**
+         * API name: {@code search_after}
+         * <p>
+         * Adds one or more values to <code>searchAfter</code>.
+         *
+         * <p><b>NOTE: In version 3.0.0 of opensearch-java, this method will be renamed to replace {@link #searchAfter(String, String...)}.</b></p>
+         */
+        public final Builder searchAfterVals(FieldValue value, FieldValue... values) {
             this.searchAfter = _listAdd(this.searchAfter, value, values);
             return this;
         }
@@ -1092,7 +1133,7 @@ protected static void setupMultisearchBodyDeserializer(ObjectDeserializer<Multis
         op.add(Builder::from, JsonpDeserializer.integerDeserializer(), "from");
         op.add(Builder::minScore, JsonpDeserializer.doubleDeserializer(), "min_score");
         op.add(Builder::postFilter, Query._DESERIALIZER, "post_filter");
-        op.add(Builder::searchAfter, JsonpDeserializer.arrayDeserializer(JsonpDeserializer.stringDeserializer()), "search_after");
+        op.add(Builder::searchAfterVals, JsonpDeserializer.arrayDeserializer(FieldValue._DESERIALIZER), "search_after");
         op.add(Builder::size, JsonpDeserializer.integerDeserializer(), "size");
         op.add(Builder::sort, JsonpDeserializer.arrayDeserializer(SortOptions._DESERIALIZER), "sort");
         op.add(Builder::trackScores, JsonpDeserializer.booleanDeserializer(), "track_scores");