Add support for project_routing for _search and _async_search

Author: pawankartik-elastic

server/src/main/java/org/elasticsearch/action/search/SearchRequest.java

@@ -116,6 +116,7 @@ public class SearchRequest extends LegacyActionRequest implements IndicesRequest
      * enabling synthetic source natively in the index.
      */
     private boolean forceSyntheticSource = false;
+    private String projectRouting;
 
     public SearchRequest() {
         this.localClusterAlias = null;
@@ -168,6 +169,15 @@ public boolean allowsCrossProject() {
         return true;
     }
 
+    @Override
+    public String getProjectRouting() {
+        return projectRouting;
+    }
+
+    public void setProjectRouting(@Nullable String projectRouting) {
+        this.projectRouting = projectRouting;
+    }
+
     /**
      * Creates a new sub-search request starting from the original search request that is provided.
      * For internal use only, allows to fork a search request into multiple search requests that will be executed independently.

server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java

@@ -620,7 +620,8 @@ public void onFailure(Exception e) {
                         }),
                         forceConnectTimeoutSecs,
                         resolvesCrossProject,
-                        rewritten.getResolvedIndexExpressions()
+                        rewritten.getResolvedIndexExpressions(),
+                        rewritten.getProjectRouting()
                     );
                 }
             }
@@ -1065,7 +1066,8 @@ static void collectSearchShards(
         ActionListener<Map<String, SearchShardsResponse>> listener,
         TimeValue forceConnectTimeoutSecs,
         boolean resolvesCrossProject,
-        ResolvedIndexExpressions originResolvedIdxExpressions
+        ResolvedIndexExpressions originResolvedIdxExpressions,
+        String projectRouting
     ) {
         RemoteClusterService remoteClusterService = transportService.getRemoteClusterService();
         final CountDown responsesCountDown = new CountDown(remoteIndicesByCluster.size());
@@ -1104,7 +1106,7 @@ Map<String, SearchShardsResponse> createFinalResponse() {
                         // We do not use the relaxed index options here when validating indices' existence.
                         ElasticsearchException validationEx = CrossProjectIndexResolutionValidator.validate(
                             originalIdxOpts,
-                            null,
+                            projectRouting,
                             originResolvedIdxExpressions,
                             resolvedIndexExpressions
                         );

server/src/main/java/org/elasticsearch/rest/action/search/RestSearchAction.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.rest.action.RestRefCountedChunkedToXContentListener;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.crossproject.CrossProjectModeDecider;
 import org.elasticsearch.search.fetch.StoredFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.internal.SearchContext;
@@ -68,15 +69,13 @@ public class RestSearchAction extends BaseRestHandler {
     private final SearchUsageHolder searchUsageHolder;
     private final Predicate<NodeFeature> clusterSupportsFeature;
     private final Settings settings;
-
-    public RestSearchAction(SearchUsageHolder searchUsageHolder, Predicate<NodeFeature> clusterSupportsFeature) {
-        this(searchUsageHolder, clusterSupportsFeature, null);
-    }
+    private final CrossProjectModeDecider crossProjectModeDecider;
 
     public RestSearchAction(SearchUsageHolder searchUsageHolder, Predicate<NodeFeature> clusterSupportsFeature, Settings settings) {
         this.searchUsageHolder = searchUsageHolder;
         this.clusterSupportsFeature = clusterSupportsFeature;
         this.settings = settings;
+        this.crossProjectModeDecider = new CrossProjectModeDecider(settings);
     }
 
     @Override
@@ -109,10 +108,9 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         // this might be set by old clients
         request.param("min_compatible_shard_node");
 
-        final boolean crossProjectEnabled = settings != null && settings.getAsBoolean("serverless.cross_project.enabled", false);
-        if (crossProjectEnabled) {
-            // accept but drop project_routing param until fully supported
-            request.param("project_routing");
+        final boolean crossProjectEnabled = crossProjectModeDecider.crossProjectEnabled();
+        if (crossProjectEnabled && searchRequest.allowsCrossProject()) {
+            searchRequest.setProjectRouting(request.param("project_routing"));
         }
 
         /*
@@ -252,7 +250,7 @@ public static void parseSearchRequest(
         }
         searchRequest.indicesOptions(indicesOptions);
 
-        validateSearchRequest(request, searchRequest);
+        validateSearchRequest(request, searchRequest, crossProjectEnabled && searchRequest.allowsCrossProject());
 
         if (searchRequest.pointInTimeBuilder() != null) {
             preparePointInTime(searchRequest, request);
@@ -412,10 +410,46 @@ static void preparePointInTime(SearchRequest request, RestRequest restRequest) {
      * might modify the search request to align certain parameters.
      */
     public static void validateSearchRequest(RestRequest restRequest, SearchRequest searchRequest) {
+        validateSearchRequest(restRequest, searchRequest, false);
+    }
+
+    private static void validateSearchRequest(RestRequest restRequest, SearchRequest searchRequest, boolean resolvesCrossProject) {
         checkRestTotalHits(restRequest, searchRequest);
         checkSearchType(restRequest, searchRequest);
         // ensures that the rest param is consumed
         restRequest.paramAsBoolean(INCLUDE_NAMED_QUERIES_SCORE_PARAM, false);
+        checkProjectRouting(searchRequest, resolvesCrossProject);
+    }
+
+    private static void checkProjectRouting(SearchRequest searchRequest, boolean resolvesCrossProject) {
+        /*
+         * There are 2 ways of specifying project_routing:
+         *   - as a query parameter: /_search?project_routing=..., and,
+         *   - within the request's body.
+         *
+         * Because we do not have access to `IndicesRequest/SearchRequest` from `SearchSourceBuilder`, and, project_routing
+         * can be potentially specified in 2 different places, we need to explicitly check this scenario.
+         */
+        String projectRoutingInBody = searchRequest.source().projectRouting();
+        // If it's null, either the query parameter is also null or it isn't. Either way, we're fine with it.
+        if (projectRoutingInBody != null) {
+            if (resolvesCrossProject == false) {
+                throw new IllegalArgumentException("project_routing is allowed only when CPS is enabled and the endpoint supports CPS");
+            }
+
+            // Query parameter was also set. This is not allowed, irrespective of the values.
+            if (searchRequest.getProjectRouting() != null) {
+                throw new IllegalArgumentException(
+                    "project_routing is specified in both the places: as query parameter and in the request's body"
+                );
+            }
+
+            /*
+             * Bring forward the project_routing value so that TransportSearchAction can pick it up. Although TSA can pick it up
+             * from `SearchSourceBuilder`, let's use `IndicesRequest#getProjectRouting()` since that's the intended way.
+             */
+            searchRequest.setProjectRouting(projectRoutingInBody);
+        }
     }
 
     /**

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -129,6 +129,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
     public static final ParseField POINT_IN_TIME = new ParseField("pit");
     public static final ParseField RUNTIME_MAPPINGS_FIELD = new ParseField("runtime_mappings");
     public static final ParseField RETRIEVER = new ParseField("retriever");
+    public static final ParseField PROJECT_ROUTING = new ParseField("project_routing");
 
     private static final boolean RANK_SUPPORTED = Booleans.parseBoolean(System.getProperty("es.search.rank_supported"), true);
 
@@ -212,6 +213,8 @@ public static HighlightBuilder highlight() {
 
     private boolean skipInnerHits = false;
 
+    private String projectRouting;
+
     /**
      * Constructs a new search source builder.
      */
@@ -609,6 +612,19 @@ public TimeValue timeout() {
         return timeout;
     }
 
+    public String projectRouting() {
+        return projectRouting;
+    }
+
+    public SearchSourceBuilder projectRouting(String projectRouting) {
+        if (this.projectRouting != null) {
+            throw new IllegalArgumentException("project_routing is already set");
+        }
+
+        this.projectRouting = projectRouting;
+        return this;
+    }
+
     /**
      * An optional terminate_after to terminate the search after collecting
      * <code>terminateAfter</code> documents
@@ -1382,7 +1398,9 @@ private SearchSourceBuilder parseXContent(
             if (token == XContentParser.Token.FIELD_NAME) {
                 currentFieldName = parser.currentName();
             } else if (token.isValue()) {
-                if (FROM_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                if (PROJECT_ROUTING.match(currentFieldName, parser.getDeprecationHandler())) {
+                    projectRouting(parser.text());
+                } else if (FROM_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     from(parser.intValue());
                 } else if (SIZE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     size(parser.intValue());

server/src/test/java/org/elasticsearch/rest/action/search/RestSearchActionTests.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchType;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.suggest.SuggestBuilder;
@@ -33,7 +34,7 @@ public final class RestSearchActionTests extends RestActionTestCase {
 
     @Before
     public void setUpAction() {
-        action = new RestSearchAction(new UsageService().getSearchUsageHolder(), nf -> false);
+        action = new RestSearchAction(new UsageService().getSearchUsageHolder(), nf -> false, Settings.EMPTY);
         controller().registerHandler(action);
         verifyingClient.setExecuteVerifier((actionType, request) -> mock(SearchResponse.class));
         verifyingClient.setExecuteLocallyVerifier((actionType, request) -> mock(SearchResponse.class));

x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/RestSubmitAsyncSearchAction.java

@@ -69,9 +69,8 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         SubmitAsyncSearchRequest submit = new SubmitAsyncSearchRequest();
 
         boolean crossProjectEnabled = crossProjectModeDecider.crossProjectEnabled();
-        if (crossProjectEnabled) {
-            // accept but drop project_routing param until fully supported
-            request.param("project_routing");
+        if (crossProjectEnabled && submit.getSearchRequest().allowsCrossProject()) {
+            submit.getSearchRequest().setProjectRouting(request.param("project_routing"));
         }
 
         IntConsumer setSize = size -> submit.getSearchRequest().source().size(size);