Enable Cross Project Search for templated searches along with project_routing support

Author: pawankartik-elastic

modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/MultiSearchTemplateRequest.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContent;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -37,6 +38,9 @@ public class MultiSearchTemplateRequest extends LegacyActionRequest implements C
 
     private IndicesOptions indicesOptions = IndicesOptions.strictExpandOpenAndForbidClosedIgnoreThrottled();
 
+    @Nullable
+    private String projectRouting;
+
     public MultiSearchTemplateRequest() {}
 
     MultiSearchTemplateRequest(StreamInput in) throws IOException {
@@ -153,4 +157,16 @@ public static byte[] writeMultiLineFormat(MultiSearchTemplateRequest multiSearch
         return output.toByteArray();
     }
 
+    public void setProjectRouting(@Nullable String projectRouting) {
+        if (this.projectRouting != null) {
+            throw new IllegalArgumentException("project_routing already set");
+        }
+
+        this.projectRouting = projectRouting;
+    }
+
+    @Nullable
+    public String getProjectRouting() {
+        return projectRouting;
+    }
 }

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

@@ -70,12 +70,12 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
      */
     public MultiSearchTemplateRequest parseRequest(RestRequest restRequest, boolean allowExplicitIndex) throws IOException {
         boolean crossProjectEnabled = crossProjectModeDecider.crossProjectEnabled();
+        MultiSearchTemplateRequest multiRequest = new MultiSearchTemplateRequest();
+
         if (crossProjectEnabled) {
-            // accept but drop project_routing param until fully supported
-            restRequest.param("project_routing");
+            multiRequest.setProjectRouting(restRequest.param("project_routing"));
         }
 
-        MultiSearchTemplateRequest multiRequest = new MultiSearchTemplateRequest();
         if (restRequest.hasParam("max_concurrent_searches")) {
             multiRequest.maxConcurrentSearchRequests(restRequest.paramAsInt("max_concurrent_searches", 0));
         }
@@ -86,6 +86,18 @@ public MultiSearchTemplateRequest parseRequest(RestRequest restRequest, boolean
             allowExplicitIndex,
             (searchRequest, bytes) -> {
                 SearchTemplateRequest searchTemplateRequest = SearchTemplateRequest.fromXContent(bytes);
+                if (crossProjectEnabled) {
+                    /*
+                     * If project_routing is set in the request body, it gets overwritten when parsing the body.
+                     * See SearchTemplateRequest#fromXContent() and SearchTemplateRequest#PROJECT_ROUTING_FIELD.
+                     */
+                    searchTemplateRequest.setProjectRouting(multiRequest.getProjectRouting());
+                } else {
+                    // project_routing is specified in the body and we're not in a CPS environment.
+                    if (searchTemplateRequest.getProjectRouting() != null) {
+                        throw new IllegalArgumentException("Unknown key for a VALUE_STRING in [project_routing]");
+                    }
+                }
                 if (searchTemplateRequest.getScript() != null) {
                     searchTemplateRequest.setRequest(searchRequest);
                     multiRequest.add(searchTemplateRequest);

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

@@ -64,10 +64,6 @@ public String getName() {
     @Override
     public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
         boolean crossProjectEnabled = crossProjectModeDecider.crossProjectEnabled();
-        if (crossProjectEnabled) {
-            // accept but drop project_routing param until fully supported
-            request.param("project_routing");
-        }
 
         // Creates the search request with all required params
         SearchRequest searchRequest = new SearchRequest();
@@ -86,6 +82,12 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
         try (XContentParser parser = request.contentOrSourceParamParser()) {
             searchTemplateRequest = SearchTemplateRequest.fromXContent(parser);
         }
+
+        if (searchRequest.indicesOptions().resolveCrossProjectIndexExpression() == false
+            && searchTemplateRequest.getProjectRouting() != null) {
+            throw new IllegalArgumentException("Unknown key for a VALUE_STRING in [project_routing]");
+        }
+
         searchTemplateRequest.setRequest(searchRequest);
         // This param is parsed within the search request
         if (searchRequest.source().explain() != null) {

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

@@ -12,11 +12,14 @@
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.CompositeIndicesRequest;
 import org.elasticsearch.action.LegacyActionRequest;
+import org.elasticsearch.action.ValidateActions;
 import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
@@ -44,6 +47,9 @@ public class SearchTemplateRequest extends LegacyActionRequest implements Compos
     private String script;
     private Map<String, Object> scriptParams;
 
+    @Nullable
+    private String projectRouting;
+
     public SearchTemplateRequest() {}
 
     public SearchTemplateRequest(StreamInput in) throws IOException {
@@ -138,6 +144,19 @@ public void setScriptParams(Map<String, Object> scriptParams) {
         this.scriptParams = scriptParams;
     }
 
+    public void setProjectRouting(@Nullable String projectRouting) {
+        if (this.projectRouting != null) {
+            throw new IllegalArgumentException("project_routing already set");
+        }
+
+        this.projectRouting = projectRouting;
+    }
+
+    @Nullable
+    public String getProjectRouting() {
+        return projectRouting;
+    }
+
     @Override
     public ActionRequestValidationException validate() {
         ActionRequestValidationException validationException = null;
@@ -169,6 +188,7 @@ public ActionRequestValidationException validate() {
     private static final ParseField PARAMS_FIELD = new ParseField("params");
     private static final ParseField EXPLAIN_FIELD = new ParseField("explain");
     private static final ParseField PROFILE_FIELD = new ParseField("profile");
+    private static final ParseField PROJECT_ROUTING_FIELD = new ParseField("project_routing");
 
     private static final ObjectParser<SearchTemplateRequest, Void> PARSER;
 
@@ -194,6 +214,7 @@ public ActionRequestValidationException validate() {
                 request.setScript(parser.text());
             }
         }, SOURCE_FIELD, ObjectParser.ValueType.OBJECT_OR_STRING);
+        PARSER.declareString(SearchTemplateRequest::setProjectRouting, PROJECT_ROUTING_FIELD);
     }
 
     public static SearchTemplateRequest fromXContent(XContentParser parser) throws IOException {

modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/TransportSearchTemplateAction.java

@@ -152,6 +152,7 @@ static SearchRequest convert(
         builder.profile(searchTemplateRequest.isProfile());
         checkRestTotalHitsAsInt(searchRequest, builder);
         searchRequest.source(builder);
+        searchRequest.setProjectRouting(searchTemplateRequest.getProjectRouting());
         return searchRequest;
     }