From a464576ad7c523f39707b2525bbfb4bde320b122 Mon Sep 17 00:00:00 2001
From: "uri.nudelman" <uriel.nudelman@offerup.com>
Date: Mon, 19 Aug 2024 17:17:11 -0400
Subject: [PATCH] Fixes #1116

Signed-off-by: uri.nudelman <uriel.nudelman@offerup.com>
---
 .../client/opensearch/core/SearchRequest.java | 44 +++++++++++++++++++
 .../opensearch/core/SearchRequestTest.java    | 11 +++++
 2 files changed, 55 insertions(+)

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
index 3b7166c1fc..e1a47ab883 100644
--- 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
@@ -110,6 +110,9 @@ public class SearchRequest extends RequestBase implements PlainJsonSerializable
     @Nullable
     private final String pipeline;
 
+    @Nullable
+    private final Time cancelAfterTimeInterval;
+
     @Nullable
     private final FieldCollapse collapse;
 
@@ -246,6 +249,7 @@ private SearchRequest(Builder builder) {
         this.ccsMinimizeRoundtrips = builder.ccsMinimizeRoundtrips;
         this.phaseTook = builder.phaseTook;
         this.pipeline = builder.pipeline;
+        this.cancelAfterTimeInterval = builder.cancelAfterTimeInterval;
         this.collapse = builder.collapse;
         this.defaultOperator = builder.defaultOperator;
         this.df = builder.df;
@@ -405,6 +409,17 @@ public final String pipeline() {
         return this.pipeline;
     }
 
+    /**
+     * 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;
+    }
+
     /**
      * API name: {@code collapse}
      */
@@ -1114,6 +1129,7 @@ public Builder toBuilder() {
             .ccsMinimizeRoundtrips(ccsMinimizeRoundtrips)
             .phaseTook(phaseTook)
             .pipeline(pipeline)
+            .cancelAfterTimeInterval(cancelAfterTimeInterval)
             .collapse(collapse)
             .defaultOperator(defaultOperator)
             .df(df)
@@ -1198,6 +1214,9 @@ public static class Builder extends ObjectBuilderBase implements ObjectBuilder<S
         @Nullable
         private String pipeline;
 
+        @Nullable
+        private Time cancelAfterTimeInterval;
+
         @Nullable
         private FieldCollapse collapse;
 
@@ -1472,6 +1491,27 @@ public final Builder pipeline(@Nullable String 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;
+        }
+
+        /**
+         * 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(Function<Time.Builder, ObjectBuilder<Time>> fn) {
+            return this.cancelAfterTimeInterval(fn.apply(new Time.Builder()).build());
+        }
+
         /**
          * API name: {@code collapse}
          */
@@ -2283,6 +2323,7 @@ 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");
@@ -2405,6 +2446,9 @@ protected static void setupSearchRequestDeserializer(ObjectDeserializer<SearchRe
             if (request.scroll != null) {
                 params.put("scroll", request.scroll._toJsonString());
             }
+            if (request.cancelAfterTimeInterval != null) {
+                params.put("cancel_after_time_interval", request.cancelAfterTimeInterval._toJsonString());
+            }
             if (request.searchType != null) {
                 params.put("search_type", request.searchType.jsonValue());
             }
diff --git a/java-client/src/test/java/org/opensearch/client/opensearch/core/SearchRequestTest.java b/java-client/src/test/java/org/opensearch/client/opensearch/core/SearchRequestTest.java
index addeb30493..e5824981f1 100644
--- a/java-client/src/test/java/org/opensearch/client/opensearch/core/SearchRequestTest.java
+++ b/java-client/src/test/java/org/opensearch/client/opensearch/core/SearchRequestTest.java
@@ -13,6 +13,7 @@
 import org.junit.Test;
 import org.opensearch.client.json.JsonData;
 import org.opensearch.client.opensearch._types.FieldValue;
+import org.opensearch.client.opensearch._types.Time;
 import org.opensearch.client.opensearch.core.search.SourceConfig;
 import org.opensearch.client.opensearch.core.search.SourceFilter;
 import org.opensearch.client.opensearch.model.ModelTestCase;
@@ -53,6 +54,16 @@ public void pipeline() {
         assertEquals("my_pipeline", SearchRequest._ENDPOINT.queryParameters(request).get("search_pipeline"));
     }
 
+    @Test
+    public void cancelAfterTimeInterval() {
+        Time cancelAfterTimeInterval = Time.of(ti -> ti.time("1000ms"));
+        SearchRequest request = new SearchRequest.Builder().cancelAfterTimeInterval(cancelAfterTimeInterval).build();
+
+        assertEquals("{}", toJson(request));
+        assertEquals(cancelAfterTimeInterval, request.cancelAfterTimeInterval());
+        assertEquals("1000ms", SearchRequest._ENDPOINT.queryParameters(request).get("cancel_after_time_interval"));
+    }
+
     @Test
     public void toBuilder() {
         SearchRequest origin = new SearchRequest.Builder().index("index").build();