Add support for search pipeline in search and msearch template

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce SecureHttpTransportParameters experimental API (to complement SecureTransportParameters counterpart) ([#18572](https://github.com/opensearch-project/OpenSearch/issues/18572))
 - Create equivalents of JSM's AccessController in the java agent ([#18346](https://github.com/opensearch-project/OpenSearch/issues/18346))
 - Introduced a new cluster-level API to fetch remote store metadata (segments and translogs) for each shard of an index. ([#18257](https://github.com/opensearch-project/OpenSearch/pull/18257))
+- Add support for search pipeline in search and msearch template ([#18564](https://github.com/opensearch-project/OpenSearch/pull/18564))
 
 ### Changed
 - Update Subject interface to use CheckedRunnable ([#18570](https://github.com/opensearch-project/OpenSearch/issues/18570))

client/rest-high-level/src/test/java/org/opensearch/client/RequestConvertersTests.java

@@ -1367,6 +1367,7 @@ public void testSearchTemplate() throws Exception {
         scriptParams.put("field", "name");
         scriptParams.put("value", "soren");
         searchTemplateRequest.setScriptParams(scriptParams);
+        searchTemplateRequest.setSearchPipeline("pipeline");
 
         // Verify that the resulting REST request looks as expected.
         Request request = RequestConverters.searchTemplate(searchTemplateRequest);
@@ -1391,6 +1392,7 @@ public void testRenderSearchTemplate() throws Exception {
         searchTemplateRequest.setScript("template1");
         searchTemplateRequest.setScriptType(ScriptType.STORED);
         searchTemplateRequest.setProfile(randomBoolean());
+        searchTemplateRequest.setSearchPipeline("pipeline");
 
         Map<String, Object> scriptParams = new HashMap<>();
         scriptParams.put("field", "name");
@@ -1431,6 +1433,7 @@ public void testMultiSearchTemplate() throws Exception {
             searchTemplateRequest.setScript("{\"query\": { \"match\" : { \"{{field}}\" : \"{{value}}\" }}}");
             searchTemplateRequest.setScriptType(ScriptType.INLINE);
             searchTemplateRequest.setProfile(randomBoolean());
+            searchTemplateRequest.setSearchPipeline("pipeline");
 
             Map<String, Object> scriptParams = new HashMap<>();
             scriptParams.put("field", "name");

modules/lang-mustache/src/internalClusterTest/java/org/opensearch/script/mustache/SearchTemplateIT.java

@@ -378,4 +378,87 @@ public void testIndexedTemplateWithArray() throws Exception {
         assertHitCount(searchResponse.getResponse(), 5);
     }
 
+    public void testMultiSearchTemplateWithSearchPipeline() throws Exception {
+        createIndex("my-nlp-index1");
+
+        // Indexing test data
+        client().prepareIndex("my-nlp-index1")
+            .setId("1")
+            .setSource(jsonBuilder().startObject().field("play_name", "zoo").endObject())
+            .get();
+        client().prepareIndex("my-nlp-index1")
+            .setId("2")
+            .setSource(jsonBuilder().startObject().field("play_name", "hello").endObject())
+            .get();
+        client().admin().indices().prepareRefresh().get();
+
+        // Put stored templates
+        assertAcked(
+            client().admin()
+                .cluster()
+                .preparePutStoredScript()
+                .setId("search_template_1")
+                .setContent(
+                    new BytesArray(
+                        "{\"script\": {\"lang\": \"mustache\", \"source\": {\"query\": {\"match\": {\"play_name\": \"{{play_name}}\"}}}}}"
+                    ),
+                    MediaTypeRegistry.JSON
+                )
+        );
+
+        assertAcked(
+            client().admin()
+                .cluster()
+                .preparePutStoredScript()
+                .setId("search_template_2")
+                .setContent(
+                    new BytesArray(
+                        "{\"script\": {\"lang\": \"mustache\", \"source\": {\"query\": {\"match\": {\"play_name\": \"{{play_name}}\"}}}}}"
+                    ),
+                    MediaTypeRegistry.JSON
+                )
+        );
+
+        // Create first template request with pipeline
+        String template1 = "{"
+            + "  \"id\" : \"search_template_1\","
+            + "  \"params\": {\"play_name\": \"zoo\"},"
+            + "  \"search_pipeline\": \"my_pipeline2\""
+            + "}";
+        SearchTemplateRequest request1 = SearchTemplateRequest.fromXContent(createParser(JsonXContent.jsonXContent, template1));
+        request1.setRequest(new SearchRequest("my-nlp-index1"));
+
+        // Create second template request with pipeline
+        String template2 = "{"
+            + "  \"id\" : \"search_template_2\","
+            + "  \"params\": {\"play_name\": \"hello\"},"
+            + "  \"search_pipeline\": \"my_pipeline1\""
+            + "}";
+        SearchTemplateRequest request2 = SearchTemplateRequest.fromXContent(createParser(JsonXContent.jsonXContent, template2));
+        request2.setRequest(new SearchRequest("my-nlp-index1"));
+
+        Map<String, Object> templateParams = new HashMap<>();
+        templateParams.put("play_name", "hello");
+
+        // Create second template with SearchTemplateRequestBuilder
+        SearchTemplateResponse response2 = new SearchTemplateRequestBuilder(client()).setRequest(new SearchRequest("my-nlp-index1"))
+            .setScript("search_template_2")
+            .setScriptType(ScriptType.STORED)
+            .setScriptParams(templateParams)
+            .setSearchPipeline("my_pipeline1")
+            .get();
+
+        // Execute both requests
+        SearchTemplateResponse response1 = client().execute(SearchTemplateAction.INSTANCE, request1).get();
+
+        assertNotNull(response1.getResponse());
+        assertNotNull(response2.getResponse());
+        assertHitCount(response1.getResponse(), 1);
+        assertHitCount(response2.getResponse(), 1);
+
+        // Verify that the pipeline names are set correctly in the request objects
+        assertEquals("my_pipeline2", request1.getSearchPipeline());
+        assertEquals("my_pipeline1", request2.getSearchPipeline());
+    }
+
 }

modules/lang-mustache/src/main/java/org/opensearch/script/mustache/RestMultiSearchTemplateAction.java

@@ -107,6 +107,11 @@ public static MultiSearchTemplateRequest parseRequest(RestRequest restRequest, b
             (searchRequest, bytes) -> {
                 SearchTemplateRequest searchTemplateRequest = SearchTemplateRequest.fromXContent(bytes);
                 if (searchTemplateRequest.getScript() != null) {
+                    // Set the search request pipeline
+                    String pipeline = searchTemplateRequest.getSearchPipeline();
+                    if (pipeline != null && !pipeline.isEmpty()) {
+                        searchRequest.pipeline(pipeline);
+                    }
                     searchTemplateRequest.setRequest(searchRequest);
                     multiRequest.add(searchTemplateRequest);
                 } else {

modules/lang-mustache/src/main/java/org/opensearch/script/mustache/RestSearchTemplateAction.java

@@ -95,6 +95,11 @@
         try (XContentParser parser = request.contentOrSourceParamParser()) {
             searchTemplateRequest = SearchTemplateRequest.fromXContent(parser);
         }
+        // Set the search request pipeline
+        String pipeline = searchTemplateRequest.getSearchPipeline();
+        if (pipeline != null && !pipeline.isEmpty()) {
+            searchRequest.pipeline(pipeline);
+        }
         searchTemplateRequest.setRequest(searchRequest);
 
         return channel -> client.execute(SearchTemplateAction.INSTANCE, searchTemplateRequest, new RestStatusToXContentListener<>(channel));

modules/lang-mustache/src/main/java/org/opensearch/script/mustache/SearchTemplateRequest.java

@@ -32,6 +32,7 @@
 
 package org.opensearch.script.mustache;
 
+import org.opensearch.Version;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.action.CompositeIndicesRequest;
@@ -67,6 +68,7 @@ public class SearchTemplateRequest extends ActionRequest implements IndicesReque
     private ScriptType scriptType;
     private String script;
     private Map<String, Object> scriptParams;
+    private String searchPipeline;
 
     public SearchTemplateRequest() {}
 
@@ -81,6 +83,10 @@ public SearchTemplateRequest(StreamInput in) throws IOException {
         if (in.readBoolean()) {
             scriptParams = in.readMap();
         }
+        if (in.getVersion().onOrAfter(Version.V_3_2_0)) {
+            searchPipeline = in.readOptionalString();
+        }
+
     }
 
     public SearchTemplateRequest(SearchRequest searchRequest) {
@@ -106,12 +112,13 @@ public boolean equals(Object o) {
             && Objects.equals(request, request1.request)
             && scriptType == request1.scriptType
             && Objects.equals(script, request1.script)
-            && Objects.equals(scriptParams, request1.scriptParams);
+            && Objects.equals(scriptParams, request1.scriptParams)
+            && Objects.equals(searchPipeline, request1.searchPipeline);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(request, simulate, explain, profile, scriptType, script, scriptParams);
+        return Objects.hash(request, simulate, explain, profile, scriptType, script, scriptParams, searchPipeline);
     }
 
     public boolean isSimulate() {
@@ -162,6 +169,14 @@ public void setScriptParams(Map<String, Object> scriptParams) {
         this.scriptParams = scriptParams;
     }
 
+    public String getSearchPipeline() {
+        return searchPipeline;
+    }
+
+    public void setSearchPipeline(String searchPipeline) {
+        this.searchPipeline = searchPipeline;
+    }
+
     @Override
     public ActionRequestValidationException validate() {
         ActionRequestValidationException validationException = null;
@@ -193,6 +208,7 @@ public ActionRequestValidationException validate() {
     private static ParseField PARAMS_FIELD = new ParseField("params");
     private static ParseField EXPLAIN_FIELD = new ParseField("explain");
     private static ParseField PROFILE_FIELD = new ParseField("profile");
+    private static ParseField SEARCH_PIPELINE_FIELD = new ParseField("search_pipeline");
 
     private static final ObjectParser<SearchTemplateRequest, Void> PARSER;
     static {
@@ -217,6 +233,14 @@ public ActionRequestValidationException validate() {
                 request.setScript(parser.text());
             }
         }, SOURCE_FIELD, ObjectParser.ValueType.OBJECT_OR_STRING);
+        PARSER.declareField((parser, request, context) -> {
+            if (parser.currentToken() == XContentParser.Token.VALUE_NULL) {
+                request.setSearchPipeline(null);
+            } else {
+                request.setSearchPipeline(parser.text());
+            }
+        }, SEARCH_PIPELINE_FIELD, ObjectParser.ValueType.STRING_OR_NULL);
+
     }
 
     public static SearchTemplateRequest fromXContent(XContentParser parser) throws IOException {
@@ -238,6 +262,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         return builder.field(PARAMS_FIELD.getPreferredName(), scriptParams)
             .field(EXPLAIN_FIELD.getPreferredName(), explain)
             .field(PROFILE_FIELD.getPreferredName(), profile)
+            .field(SEARCH_PIPELINE_FIELD.getPreferredName(), searchPipeline)
             .endObject();
     }
 
@@ -255,6 +280,9 @@ public void writeTo(StreamOutput out) throws IOException {
         if (hasParams) {
             out.writeMap(scriptParams);
         }
+        if (out.getVersion().onOrAfter(Version.V_3_2_0)) {
+            out.writeOptionalString(searchPipeline);
+        }
     }
 
     @Override

modules/lang-mustache/src/main/java/org/opensearch/script/mustache/SearchTemplateRequestBuilder.java

@@ -89,4 +89,9 @@
         request.setScriptParams(scriptParams);
         return this;
     }
+
+    public SearchTemplateRequestBuilder setSearchPipeline(String searchPipeline) {
+        request.setSearchPipeline(searchPipeline);
+        return this;
+    }
 }

modules/lang-mustache/src/test/java/org/opensearch/script/mustache/MultiSearchTemplateRequestTests.java

@@ -126,13 +126,15 @@ public void testMultiSearchTemplateToJson() throws Exception {
             SearchRequest searchRequest = new SearchRequest(indices);
             // scroll is not supported in the current msearch or msearchtemplate api, so unset it:
             searchRequest.scroll((Scroll) null);
+            searchRequest.pipeline("pipeline");
             // batched reduce size is currently not set-able on a per-request basis as it is a query string parameter only
             searchRequest.setBatchedReduceSize(SearchRequest.DEFAULT_BATCHED_REDUCE_SIZE);
             SearchTemplateRequest searchTemplateRequest = new SearchTemplateRequest(searchRequest);
 
             searchTemplateRequest.setScript("{\"query\": { \"match\" : { \"{{field}}\" : \"{{value}}\" }}}");
             searchTemplateRequest.setScriptType(ScriptType.INLINE);
             searchTemplateRequest.setProfile(randomBoolean());
+            searchTemplateRequest.setSearchPipeline("pipeline");
 
             Map<String, Object> scriptParams = new HashMap<>();
             scriptParams.put("field", "name");
@@ -166,6 +168,34 @@ public void testMultiSearchTemplateToJson() throws Exception {
         assertEquals(serialized, toJsonString(deser));
     }
 
+    public void testParseRequestWithSearchPipeline() throws Exception {
+        byte[] data = StreamsUtils.copyToBytesFromClasspath("/org/opensearch/script/mustache/simple-msearch-template.json");
+        RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withContent(new BytesArray(data), MediaTypeRegistry.JSON)
+            .build();
+
+        MultiSearchTemplateRequest request = RestMultiSearchTemplateAction.parseRequest(restRequest, true);
+
+        assertThat(request.requests().size(), equalTo(3));
+        SearchTemplateRequest searchTemplateRequest = request.requests().get(0);
+
+        assertThat(request.requests().get(0).getRequest().indices()[0], equalTo("test0"));
+        assertThat(request.requests().get(0).getRequest().indices()[1], equalTo("test1"));
+        assertThat(request.requests().get(0).getRequest().indices(), arrayContaining("test0", "test1"));
+        assertThat(request.requests().get(0).getRequest().pipeline(), equalTo("my_pipeline"));
+        assertThat(request.requests().get(1).getRequest().indices()[0], equalTo("test2"));
+        assertThat(request.requests().get(1).getRequest().indices()[1], equalTo("test3"));
+        assertThat(request.requests().get(1).getRequest().indices(), arrayContaining("test2", "test3"));
+        assertThat(request.requests().get(1).getRequest().pipeline(), equalTo("my_pipeline1"));
+        assertThat(request.requests().get(2).getRequest().indices()[0], equalTo("test4"));
+        assertThat(request.requests().get(2).getRequest().indices()[1], equalTo("test1"));
+        assertThat(request.requests().get(2).getRequest().indices(), arrayContaining("test4", "test1"));
+        assertThat(request.requests().get(2).getRequest().pipeline(), equalTo("my_pipeline2"));
+
+        // Additional validation
+        assertEquals("{\"query\":{\"match_{{template}}\":{}}}", searchTemplateRequest.getScript());
+        assertEquals("all", searchTemplateRequest.getScriptParams().get("template"));
+    }
+
     protected String toJsonString(MultiSearchTemplateRequest multiSearchTemplateRequest) throws IOException {
         byte[] bytes = MultiSearchTemplateRequest.writeMultiLineFormat(multiSearchTemplateRequest, MediaTypeRegistry.JSON.xContent());
         return new String(bytes, StandardCharsets.UTF_8);

modules/lang-mustache/src/test/java/org/opensearch/script/mustache/SearchTemplateRequestXContentTests.java

@@ -94,6 +94,7 @@ public void testToXContentWithInlineTemplate() throws IOException {
         request.setScriptType(ScriptType.INLINE);
         request.setScript("{\"query\": { \"match\" : { \"{{my_field}}\" : \"{{my_value}}\" } } }");
         request.setProfile(true);
+        request.setSearchPipeline("pipeline");
 
         Map<String, Object> scriptParams = new HashMap<>();
         scriptParams.put("my_field", "foo");
@@ -110,6 +111,7 @@ public void testToXContentWithInlineTemplate() throws IOException {
             .endObject()
             .field("explain", false)
             .field("profile", true)
+            .field("search_pipeline", "pipeline")
             .endObject();
 
         XContentBuilder actualRequest = MediaTypeRegistry.contentBuilder(contentType);
@@ -124,6 +126,7 @@ public void testToXContentWithStoredTemplate() throws IOException {
         request.setScriptType(ScriptType.STORED);
         request.setScript("match_template");
         request.setExplain(true);
+        request.setSearchPipeline("pipeline");
 
         Map<String, Object> params = new HashMap<>();
         params.put("my_field", "foo");
@@ -140,6 +143,7 @@ public void testToXContentWithStoredTemplate() throws IOException {
             .endObject()
             .field("explain", true)
             .field("profile", false)
+            .field("search_pipeline", "pipeline")
             .endObject();
 
         XContentBuilder actualRequest = MediaTypeRegistry.contentBuilder(contentType);

modules/lang-mustache/src/test/resources/org/opensearch/script/mustache/simple-msearch-template.json

@@ -1,6 +1,6 @@
 {"index":["test0", "test1"], "request_cache": true}
-{"source": {"query" : {"match_{{template}}" :{}}}, "params": {"template": "all" } }
+{"source": {"query" : {"match_{{template}}" :{}}}, "params": {"template": "all" }, "search_pipeline": "my_pipeline" }
 {"index" : "test2,test3", "preference": "_local"}
-{"source": {"query" : {"match_{{template}}" :{}}}, "params": {"template": "all" } }
+{"source": {"query" : {"match_{{template}}" :{}}}, "params": {"template": "all" }, "search_pipeline": "my_pipeline1" }
 {"index" : ["test4", "test1"], "routing": "123"}
-{"source": {"query" : {"match_{{template}}" :{}}}, "params": {"template": "all" } }
+{"source": {"query" : {"match_{{template}}" :{}}}, "params": {"template": "all" }, "search_pipeline": "my_pipeline2" }

server/src/main/java/org/opensearch/action/search/MultiSearchRequest.java

@@ -222,7 +222,7 @@ public static void readMultiLineFormat(
             }
 
             SearchRequest searchRequest = new SearchRequest();
-            if (indices != null) {
+            if (indices != null && indices.length > 0) {
                 searchRequest.indices(indices);
             }
             if (indicesOptions != null) {