Simplify cluster state management in ESQL (#131635)

Author: idegtiarenko

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

@@ -8,9 +8,9 @@
 package org.elasticsearch.xpack.esql.session;
 
 import org.elasticsearch.ElasticsearchSecurityException;
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.action.OriginalIndices;
 import org.elasticsearch.action.fieldcaps.FieldCapabilitiesFailure;
 import org.elasticsearch.action.search.ShardSearchFailure;
 import org.elasticsearch.action.support.IndicesOptions;
@@ -313,48 +313,53 @@ static void updateExecutionInfoAtEndOfPlanning(EsqlExecutionInfo execInfo) {
     }
 
     /**
-     * Checks the index expression for the presence of remote clusters. If found, it will ensure that the caller
-     * has a valid Enterprise (or Trial) license on the querying cluster.
-     * @param indices index expression requested by user
-     * @param indicesGrouper grouper of index expressions by cluster alias
-     * @param licenseState license state on the querying cluster
+     * Checks the index expression for the presence of remote clusters.
+     * If found, it will ensure that the caller has a valid Enterprise (or Trial) license on the querying cluster
+     * as well as initialize corresponding cluster state in execution info.
      * @throws org.elasticsearch.ElasticsearchStatusException if the license is not valid (or present) for ES|QL CCS search.
      */
-    public static void checkForCcsLicense(
-        EsqlExecutionInfo executionInfo,
-        List<IndexPattern> indices,
+    public static void initCrossClusterState(
         IndicesExpressionGrouper indicesGrouper,
-        Set<String> configuredClusters,
-        XPackLicenseState licenseState
-    ) {
-        for (IndexPattern index : indices) {
-            Map<String, OriginalIndices> groupedIndices;
-            try {
-                groupedIndices = indicesGrouper.groupIndices(configuredClusters, IndicesOptions.DEFAULT, index.indexPattern());
-            } catch (NoSuchRemoteClusterException e) {
-                if (EsqlLicenseChecker.isCcsAllowed(licenseState)) {
-                    throw e;
-                } else {
-                    throw EsqlLicenseChecker.invalidLicenseForCcsException(licenseState);
-                }
+        XPackLicenseState licenseState,
+        List<IndexPattern> patterns,
+        EsqlExecutionInfo executionInfo
+    ) throws ElasticsearchStatusException {
+        if (patterns.isEmpty()) {
+            return;
+        }
+        assert patterns.size() == 1 : "Only single index pattern is supported";
+        try {
+            var groupedIndices = indicesGrouper.groupIndices(
+                // indicesGrouper.getConfiguredClusters() might return mutable set that changes as clusters connect or disconnect.
+                // it is copied here so that we have the same resolution when request contains multiple remote cluster patterns with *
+                Set.copyOf(indicesGrouper.getConfiguredClusters()),
+                IndicesOptions.DEFAULT,
+                patterns.getFirst().indexPattern()
+            );
+
+            // initialize the cluster entries in EsqlExecutionInfo before throwing the invalid license error
+            // so that the CCS telemetry handler can recognize that this error is CCS-related
+            for (var entry : groupedIndices.entrySet()) {
+                final String clusterAlias = entry.getKey();
+                final String indexExpr = Strings.arrayToCommaDelimitedString(entry.getValue().indices());
+                executionInfo.swapCluster(clusterAlias, (k, v) -> {
+                    assert v == null : "No cluster for " + clusterAlias + " should have been added to ExecutionInfo yet";
+                    return new EsqlExecutionInfo.Cluster(clusterAlias, indexExpr, executionInfo.isSkipUnavailable(clusterAlias));
+                });
             }
+
             // check if it is a cross-cluster query
             if (groupedIndices.size() > 1 || groupedIndices.containsKey(RemoteClusterService.LOCAL_CLUSTER_GROUP_KEY) == false) {
                 if (EsqlLicenseChecker.isCcsAllowed(licenseState) == false) {
-                    // initialize the cluster entries in EsqlExecutionInfo before throwing the invalid license error
-                    // so that the CCS telemetry handler can recognize that this error is CCS-related
-                    for (Map.Entry<String, OriginalIndices> entry : groupedIndices.entrySet()) {
-                        executionInfo.swapCluster(
-                            entry.getKey(),
-                            (k, v) -> new EsqlExecutionInfo.Cluster(
-                                entry.getKey(),
-                                Strings.arrayToCommaDelimitedString(entry.getValue().indices())
-                            )
-                        );
-                    }
                     throw EsqlLicenseChecker.invalidLicenseForCcsException(licenseState);
                 }
             }
+        } catch (NoSuchRemoteClusterException e) {
+            if (EsqlLicenseChecker.isCcsAllowed(licenseState)) {
+                throw e;
+            } else {
+                throw EsqlLicenseChecker.invalidLicenseForCcsException(licenseState);
+            }
         }
     }
 

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

@@ -11,12 +11,9 @@
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.action.OriginalIndices;
 import org.elasticsearch.action.fieldcaps.FieldCapabilitiesFailure;
 import org.elasticsearch.action.search.ShardSearchFailure;
-import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.action.support.SubscribableListener;
-import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.common.regex.Regex;
@@ -156,7 +153,6 @@ public interface PlanRunner {
     private final PhysicalPlanOptimizer physicalPlanOptimizer;
     private final PlanTelemetry planTelemetry;
     private final IndicesExpressionGrouper indicesExpressionGrouper;
-    private Set<String> configuredClusters;
     private final InferenceRunner inferenceRunner;
     private final RemoteClusterService remoteClusterService;
 
@@ -410,8 +406,6 @@ public void analyzedPlan(
             plan.setAnalyzed();
             return plan;
         };
-        // Capture configured remotes list to ensure consistency throughout the session
-        configuredClusters = Set.copyOf(indicesExpressionGrouper.getConfiguredClusters());
 
         PreAnalyzer.PreAnalysis preAnalysis = preAnalyzer.preAnalyze(parsed);
         var unresolvedPolicies = preAnalysis.enriches.stream()
@@ -422,10 +416,8 @@ public void analyzedPlan(
                 )
             )
             .collect(Collectors.toSet());
-        final List<IndexPattern> indices = preAnalysis.indices;
 
-        EsqlCCSUtils.checkForCcsLicense(executionInfo, indices, indicesExpressionGrouper, configuredClusters, verifier.licenseState());
-        initializeClusterData(indices, executionInfo);
+        EsqlCCSUtils.initCrossClusterState(indicesExpressionGrouper, verifier.licenseState(), preAnalysis.indices, executionInfo);
 
         var listener = SubscribableListener.<EnrichResolution>newForked(
             l -> enrichPolicyResolver.resolvePolicies(unresolvedPolicies, executionInfo, l)
@@ -660,26 +652,6 @@ private void validateRemoteVersions(EsqlExecutionInfo executionInfo) {
         });
     }
 
-    private void initializeClusterData(List<IndexPattern> indices, EsqlExecutionInfo executionInfo) {
-        if (indices.isEmpty()) {
-            return;
-        }
-        assert indices.size() == 1 : "Only single index pattern is supported";
-        Map<String, OriginalIndices> clusterIndices = indicesExpressionGrouper.groupIndices(
-            configuredClusters,
-            IndicesOptions.DEFAULT,
-            indices.getFirst().indexPattern()
-        );
-        for (Map.Entry<String, OriginalIndices> entry : clusterIndices.entrySet()) {
-            final String clusterAlias = entry.getKey();
-            String indexExpr = Strings.arrayToCommaDelimitedString(entry.getValue().indices());
-            executionInfo.swapCluster(clusterAlias, (k, v) -> {
-                assert v == null : "No cluster for " + clusterAlias + " should have been added to ExecutionInfo yet";
-                return new EsqlExecutionInfo.Cluster(clusterAlias, indexExpr, executionInfo.isSkipUnavailable(clusterAlias));
-            });
-        }
-    }
-
     private void preAnalyzeMainIndices(
         PreAnalyzer.PreAnalysis preAnalysis,
         EsqlExecutionInfo executionInfo,