resolve remotes with field caps call

Author: idegtiarenko

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolver.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.common.util.iterable.Iterables;
 import org.elasticsearch.core.Tuple;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.AbstractTransportRequest;
@@ -39,21 +40,24 @@
 import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
 import org.elasticsearch.xpack.esql.action.EsqlExecutionInfo;
 import org.elasticsearch.xpack.esql.analysis.EnrichResolution;
+import org.elasticsearch.xpack.esql.analysis.PreAnalyzer;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
 import org.elasticsearch.xpack.esql.core.util.StringUtils;
 import org.elasticsearch.xpack.esql.index.EsIndex;
+import org.elasticsearch.xpack.esql.index.MappingException;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput;
+import org.elasticsearch.xpack.esql.plan.IndexPattern;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
+import org.elasticsearch.xpack.esql.session.EsqlCCSUtils;
 import org.elasticsearch.xpack.esql.session.IndexResolver;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -110,22 +114,37 @@ public static UnresolvedPolicy from(Enrich e) {
     /**
      * Resolves a set of enrich policies
      *
-     * @param enriches           the unresolved policies
+     * @param preAnalysis        to retrieve indices and enriches to resolve
+     * @param requestFilter      to resolve target clusters
      * @param executionInfo      the execution info
      * @param listener           notified with the enrich resolution
      */
-    public void resolvePolicies(List<Enrich> enriches, EsqlExecutionInfo executionInfo, ActionListener<EnrichResolution> listener) {
-        if (enriches.isEmpty()) {
+    public void resolvePolicies(
+        PreAnalyzer.PreAnalysis preAnalysis,
+        QueryBuilder requestFilter,
+        EsqlExecutionInfo executionInfo,
+        ActionListener<EnrichResolution> listener
+    ) {
+        if (preAnalysis.enriches.isEmpty()) {
             listener.onResponse(new EnrichResolution());
             return;
         }
 
-        doResolvePolicies(
-            new HashSet<>(executionInfo.getClusters().keySet()),
-            enriches.stream().map(EnrichPolicyResolver.UnresolvedPolicy::from).toList(),
-            executionInfo,
-            listener
-        );
+        doResolveRemotes(preAnalysis.indices, requestFilter, listener.delegateFailureAndWrap((l, remotes) -> {
+            doResolvePolicies(remotes, preAnalysis.enriches.stream().map(UnresolvedPolicy::from).toList(), executionInfo, l);
+        }));
+    }
+
+    private void doResolveRemotes(List<IndexPattern> indexPatterns, QueryBuilder requestFilter, ActionListener<Set<String>> listener) {
+        switch (indexPatterns.size()) {
+            case 0 -> listener.onResponse(Set.of());
+            case 1 -> indexResolver.resolveConcreteIndices(
+                indexPatterns.getFirst().indexPattern(),
+                requestFilter,
+                listener.map(EsqlCCSUtils::getRemotesOf)
+            );
+            default -> listener.onFailure(new MappingException("Queries with multiple indices are not supported"));
+        }
     }
 
     protected void doResolvePolicies(
@@ -442,7 +461,7 @@ public void messageReceived(LookupRequest request, TransportChannel channel, Tas
                     if (p == null) {
                         continue;
                     }
-                    try (ThreadContext.StoredContext ignored = threadContext.stashWithOrigin(ClientHelper.ENRICH_ORIGIN)) {
+                    try (var ignored = threadContext.stashWithOrigin(ClientHelper.ENRICH_ORIGIN)) {
                         String indexName = EnrichPolicy.getBaseName(policyName);
                         indexResolver.resolveAsMergedMapping(indexName, IndexResolver.ALL_FIELDS, null, false, refs.acquire(indexResult -> {
                             if (indexResult.isValid() && indexResult.get().concreteIndices().size() == 1) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlCCSUtils.java

@@ -40,7 +40,8 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.stream.Collectors;
+
+import static java.util.stream.Collectors.toSet;
 
 public class EsqlCCSUtils {
 
@@ -206,7 +207,7 @@ static void updateExecutionInfoWithClustersWithNoMatchingIndices(
         // NOTE: we assume that updateExecutionInfoWithUnavailableClusters() was already run and took care of unavailable clusters.
         final Set<String> clustersWithNoMatchingIndices = executionInfo.getClusterStates(Cluster.Status.RUNNING)
             .map(Cluster::getClusterAlias)
-            .collect(Collectors.toSet());
+            .collect(toSet());
         for (String indexName : indexResolution.resolvedIndices()) {
             clustersWithNoMatchingIndices.remove(RemoteClusterAware.parseClusterAlias(indexName));
         }
@@ -414,4 +415,8 @@ public static String inClusterName(String clusterAlias) {
             return "in remote cluster [" + clusterAlias + "]";
         }
     }
+
+    public static Set<String> getRemotesOf(Set<String> concreteIndices) {
+        return concreteIndices.stream().map(RemoteClusterAware::parseClusterAlias).collect(toSet());
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -376,7 +376,7 @@ public void analyzedPlan(
         EsqlCCSUtils.initCrossClusterState(indicesExpressionGrouper, verifier.licenseState(), preAnalysis.indices, executionInfo);
 
         var listener = SubscribableListener. //
-        <EnrichResolution>newForked(l -> enrichPolicyResolver.resolvePolicies(preAnalysis.enriches, executionInfo, l))
+        <EnrichResolution>newForked(l -> enrichPolicyResolver.resolvePolicies(preAnalysis, requestFilter, executionInfo, l))
             .<PreAnalysisResult>andThenApply(enrichResolution -> FieldNameUtils.resolveFieldNames(parsed, enrichResolution))
             .<PreAnalysisResult>andThen((l, preAnalysisResult) -> resolveInferences(parsed, preAnalysisResult, l));
         // first resolve the lookup indices, then the main indices

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/IndexResolver.java

@@ -74,6 +74,14 @@ public IndexResolver(Client client) {
         this.client = client;
     }
 
+    public void resolveConcreteIndices(String indexPattern, QueryBuilder requestFilter, ActionListener<Set<String>> listener) {
+        client.execute(
+            EsqlResolveFieldsAction.TYPE,
+            createFieldCapsRequest(indexPattern, Set.of("_id"), requestFilter, false),
+            listener.delegateFailureAndWrap((l, response) -> l.onResponse(Set.of(response.getIndices())))
+        );
+    }
+
     /**
      * Resolves a pattern to one (potentially compound meaning that spawns multiple indices) mapping.
      */