Adds MsearchRequest. Fixes SearchRequest.

Signed-off-by: uri.nudelman <[email protected]>
opensearch-project · Aug 19, 2024 · e537648 · e537648
1 parent a464576
commit e537648
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 3 deletions.
diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/core/MsearchRequest.java b/java-client/src/main/java/org/opensearch/client/opensearch/core/MsearchRequest.java
@@ -48,6 +48,7 @@
 import org.opensearch.client.opensearch._types.ExpandWildcard;
 import org.opensearch.client.opensearch._types.RequestBase;
 import org.opensearch.client.opensearch._types.SearchType;
+import org.opensearch.client.opensearch._types.Time;
 import org.opensearch.client.opensearch.core.msearch.RequestItem;
 import org.opensearch.client.transport.Endpoint;
 import org.opensearch.client.transport.endpoints.SimpleEndpoint;
@@ -66,6 +67,9 @@ public class MsearchRequest extends RequestBase implements NdJsonpSerializable,
     @Nullable
     private final Boolean allowNoIndices;
 
+    @Nullable
+    private final Time cancelAfterTimeInterval;
+
     @Nullable
     private final Boolean ccsMinimizeRoundtrips;
 
@@ -98,6 +102,7 @@ public class MsearchRequest extends RequestBase implements NdJsonpSerializable,
     private MsearchRequest(Builder builder) {
 
         this.allowNoIndices = builder.allowNoIndices;
+        this.cancelAfterTimeInterval = builder.cancelAfterTimeInterval;
         this.ccsMinimizeRoundtrips = builder.ccsMinimizeRoundtrips;
         this.expandWildcards = ApiTypeHelper.unmodifiable(builder.expandWildcards);
         this.ignoreThrottled = builder.ignoreThrottled;
@@ -134,6 +139,17 @@ public final Boolean allowNoIndices() {
         return this.allowNoIndices;
     }
 
+    /**
+     * The time after which the search request will be canceled.
+     * Request-level parameter takes precedence over cancel_after_time_interval cluster setting.
+     * <p>
+     * API name: {@code cancel_after_time_interval}
+     */
+    @Nullable
+    public final Time cancelAfterTimeInterval() {
+        return this.cancelAfterTimeInterval;
+    }
+
     /**
      * If true, network roundtrips between the coordinating node and remote clusters
      * are minimized for cross-cluster search requests.
@@ -256,6 +272,7 @@ public void serialize(JsonGenerator generator, JsonpMapper mapper) {
 
     public Builder toBuilder() {
         return new Builder().allowNoIndices(allowNoIndices)
+            .cancelAfterTimeInterval(cancelAfterTimeInterval)
             .ccsMinimizeRoundtrips(ccsMinimizeRoundtrips)
             .expandWildcards(expandWildcards)
             .ignoreThrottled(ignoreThrottled)
@@ -278,6 +295,9 @@ public static class Builder extends ObjectBuilderBase implements ObjectBuilder<M
         @Nullable
         private Boolean allowNoIndices;
 
+        @Nullable
+        private Time cancelAfterTimeInterval;
+
         @Nullable
         private Boolean ccsMinimizeRoundtrips;
 
@@ -321,6 +341,17 @@ public final Builder allowNoIndices(@Nullable Boolean value) {
             return this;
         }
 
+        /**
+         * The time after which the search request will be canceled.
+         * Request-level parameter takes precedence over cancel_after_time_interval cluster setting.
+         * <p>
+         * API name: {@code cancel_after_time_interval}
+         */
+        public final Builder cancelAfterTimeInterval(@Nullable Time value) {
+            this.cancelAfterTimeInterval = value;
+            return this;
+        }
+
         /**
          * If true, network roundtrips between the coordinating node and remote clusters
          * are minimized for cross-cluster search requests.
@@ -540,6 +571,9 @@ public MsearchRequest build() {
         request -> {
             Map<String, String> params = new HashMap<>();
             params.put("typed_keys", "true");
+            if (request.cancelAfterTimeInterval != null) {
+                params.put("cancel_after_time_interval", request.cancelAfterTimeInterval._toJsonString());
+            }
             if (request.preFilterShardSize != null) {
                 params.put("pre_filter_shard_size", String.valueOf(request.preFilterShardSize));
             }

diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/core/SearchRequest.java b/java-client/src/main/java/org/opensearch/client/opensearch/core/SearchRequest.java
@@ -2323,7 +2323,6 @@ protected static void setupSearchRequestDeserializer(ObjectDeserializer<SearchRe
 
         op.add(Builder::source, SourceConfig._DESERIALIZER, "_source");
         op.add(Builder::aggregations, JsonpDeserializer.stringMapDeserializer(Aggregation._DESERIALIZER), "aggregations", "aggs");
-        op.add(Builder::cancelAfterTimeInterval, Time._DESERIALIZER, "cancel_after_time_interval");
         op.add(Builder::collapse, FieldCollapse._DESERIALIZER, "collapse");
         op.add(Builder::docvalueFields, JsonpDeserializer.arrayDeserializer(FieldAndFormat._DESERIALIZER), "docvalue_fields");
         op.add(Builder::explain, JsonpDeserializer.booleanDeserializer(), "explain");

diff --git a/java-client/src/test/java/org/opensearch/client/opensearch/core/MsearchRequestTest.java b/java-client/src/test/java/org/opensearch/client/opensearch/core/MsearchRequestTest.java
@@ -9,10 +9,11 @@
 package org.opensearch.client.opensearch.core;
 
 import java.util.Collections;
-import org.junit.Assert;
 import org.junit.Test;
+import org.opensearch.client.opensearch._types.Time;
+import org.opensearch.client.opensearch.model.ModelTestCase;
 
-public class MsearchRequestTest extends Assert {
+public class MsearchRequestTest extends ModelTestCase {
 
     @Test
     public void toBuilder() {
@@ -21,4 +22,17 @@ public void toBuilder() {
 
         assertEquals(copied.index(), origin.index());
     }
+
+    @Test
+    public void cancelAfterTimeInterval() {
+        Time cancelAfterTimeInterval = Time.of(ti -> ti.time("1000ms"));
+        MsearchRequest request = new MsearchRequest.Builder()
+                                        .index("index")
+                                        .searches(Collections.emptyList())
+                                        .cancelAfterTimeInterval(cancelAfterTimeInterval).build();
+
+        assertEquals("[]", toJson(request));
+        assertEquals(cancelAfterTimeInterval, request.cancelAfterTimeInterval());
+        assertEquals("1000ms", MsearchRequest._ENDPOINT.queryParameters(request).get("cancel_after_time_interval"));
+    }
 }